feat: add header factory

server/http.go

@@ -5,16 +5,26 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"net/http"
+	"regexp"
+	"strconv"
 	"strings"
+	"time"
 	"tunnel_pls/session"
 	"tunnel_pls/utils"
 
 	"github.com/gorilla/websocket"
+	"golang.org/x/crypto/ssh"
 )
 
+var BAD_GATEWAY_RESPONSE = []byte("HTTP/1.1 502 Bad Gateway\r\n" +
+	"Content-Length: 11\r\n" +
+	"Content-Type: text/plain\r\n\r\n" +
+	"Bad Gateway")
+
 var upgrader = websocket.Upgrader{
 	ReadBufferSize:  1024,
 	WriteBufferSize: 1024,
@@ -70,6 +80,172 @@ func (w *connResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 	return w.conn, rw, nil
 }
 
+type CustomWriter struct {
+	RemoteAddr  net.Addr
+	writer      io.Writer
+	reader      io.Reader
+	headerBuf   []byte
+	buf         []byte
+	Requests    []*RequestContext
+	interaction *session.Interaction
+}
+
+type RequestContext struct {
+	Host        string
+	Path        string
+	Method      string
+	Chunked     bool
+	Tail        []byte
+	ContentSize int
+	Written     int
+}
+
+func (cw *CustomWriter) Read(p []byte) (int, error) {
+	read, err := cw.reader.Read(p)
+	test := bytes.NewReader(p)
+	reqhf, _ := NewRequestHeaderFactory(test)
+	if reqhf != nil {
+		cw.Requests = append(cw.Requests, &RequestContext{
+			Host:        reqhf.Get("Host"),
+			Path:        reqhf.Path,
+			Method:      reqhf.Method,
+			Chunked:     false,
+			Tail:        make([]byte, 5),
+			ContentSize: 0,
+			Written:     0,
+		})
+	}
+	return read, err
+}
+
+func NewCustomWriter(writer io.Writer, reader io.Reader, remoteAddr net.Addr) *CustomWriter {
+	return &CustomWriter{
+		RemoteAddr:  remoteAddr,
+		writer:      writer,
+		reader:      reader,
+		buf:         make([]byte, 0, 4096),
+		interaction: nil,
+	}
+}
+
+var DELIMITER = []byte{0x0D, 0x0A, 0x0D, 0x0A} // HTTP HEADER DELIMITER `\r\n\r\n`
+var requestLine = regexp.MustCompile(`^(GET|POST|PUT|DELETE|HEAD|OPTIONS|PATCH|TRACE|CONNECT) \S+ HTTP/\d\.\d$`)
+var responseLine = regexp.MustCompile(`^HTTP/\d\.\d \d{3} .+`)
+
+func isHTTPHeader(buf []byte) bool {
+	lines := bytes.Split(buf, []byte("\r\n"))
+	if len(lines) < 1 {
+		return false
+	}
+	startLine := string(lines[0])
+	if !requestLine.MatchString(startLine) && !responseLine.MatchString(startLine) {
+		return false
+	}
+
+	for _, line := range lines[1:] {
+		if len(line) == 0 {
+			break
+		}
+		if !bytes.Contains(line, []byte(":")) {
+			return false
+		}
+	}
+	return true
+}
+
+func (cw *CustomWriter) Write(p []byte) (int, error) {
+	if len(p) == len(BAD_GATEWAY_RESPONSE) && bytes.Equal(p, BAD_GATEWAY_RESPONSE) {
+		return cw.writer.Write(p)
+	}
+
+	cw.buf = append(cw.buf, p...)
+	timestamp := time.Now().UTC().Format(time.RFC3339)
+	// TODO: implement middleware buat cache system dll
+	if idx := bytes.Index(cw.buf, DELIMITER); idx != -1 {
+		header := cw.buf[:idx+len(DELIMITER)]
+		body := cw.buf[idx+len(DELIMITER):]
+
+		if isHTTPHeader(header) {
+			resphf := NewResponseHeaderFactory(header)
+			resphf.Set("Server", "Tunnel Please")
+
+			if resphf.Get("Transfer-Encoding") == "chunked" {
+				cw.Requests[0].Chunked = true
+			}
+			if resphf.Get("Content-Length") != "" {
+				bodySize, err := strconv.Atoi(resphf.Get("Content-Length"))
+				if err != nil {
+					log.Printf("Error parsing Content-Length: %v", err)
+					cw.Requests[0].ContentSize = -1
+				} else {
+					cw.Requests[0].ContentSize = bodySize
+				}
+			} else {
+				cw.Requests[0].ContentSize = -1
+			}
+
+			header = resphf.Finalize()
+			_, err := cw.writer.Write(header)
+			if err != nil {
+				return 0, err
+			}
+
+			if len(body) > 0 {
+				_, err := cw.writer.Write(body)
+				if err != nil {
+					return 0, err
+				}
+			}
+			req := cw.Requests[0]
+			req.Written += len(body)
+
+			if req.Chunked {
+				req.Tail = append(req.Tail, p[len(p)-5:]...)
+				if bytes.Equal(req.Tail, []byte("0\r\n\r\n")) {
+					cw.interaction.SendMessage(fmt.Sprintf("\033[32m%s %s -> %s %s \033[0m\r\n", timestamp, cw.RemoteAddr.String(), req.Method, req.Path))
+				}
+			} else if req.ContentSize != -1 {
+				if req.Written >= req.ContentSize {
+					cw.Requests = cw.Requests[1:]
+					cw.interaction.SendMessage(fmt.Sprintf("\033[32m%s %s -> %s %s \033[0m\r\n", timestamp, cw.RemoteAddr.String(), req.Method, req.Path))
+				}
+			} else {
+				cw.Requests = cw.Requests[1:]
+			}
+			cw.buf = nil
+			return len(p), nil
+		}
+	}
+	cw.buf = nil
+	n, err := cw.writer.Write(p)
+	if err != nil {
+		return n, err
+	}
+
+	req := cw.Requests[0]
+	req.Written += len(p)
+	if req.Chunked {
+		req.Tail = append(req.Tail, p[len(p)-5:]...)
+		if bytes.Equal(req.Tail, []byte("0\r\n\r\n")) {
+			cw.Requests = cw.Requests[1:]
+			cw.interaction.SendMessage(fmt.Sprintf("\033[32m%s %s -> %s %s \033[0m\r\n", timestamp, cw.RemoteAddr.String(), req.Method, req.Path))
+		}
+	} else if req.ContentSize != -1 {
+		if req.Written >= req.ContentSize {
+			cw.Requests = cw.Requests[1:]
+			cw.interaction.SendMessage(fmt.Sprintf("\033[32m%s %s -> %s %s \033[0m\r\n", timestamp, cw.RemoteAddr.String(), req.Method, req.Path))
+		}
+	} else {
+		cw.Requests = cw.Requests[1:]
+	}
+
+	return n, nil
+}
+
+func (cw *CustomWriter) AddInteraction(interaction *session.Interaction) {
+	cw.interaction = interaction
+}
+
 var redirectTLS = false
 var allowedCors = make(map[string]bool)
 var isAllowedAllCors = false
@@ -123,14 +299,30 @@ func NewHTTPServer() error {
 }
 
 func Handler(conn net.Conn) {
-	reader := bufio.NewReader(conn)
-	headers, err := peekUntilHeaders(reader, 8192)
+	defer func() {
+		err := conn.Close()
+		if err != nil {
+			log.Printf("Error closing connection: %v", err)
+			return
+		}
+		return
+	}()
+
+	dstReader := bufio.NewReader(conn)
+	reqhf, err := NewRequestHeaderFactory(dstReader)
 	if err != nil {
-		log.Println("Failed to peek headers:", err)
 		return
 	}
+	cw := NewCustomWriter(conn, dstReader, conn.RemoteAddr())
 
-	host := strings.Split(parseHostFromHeader(headers), ".")
+	// Initial Requests
+	cw.Requests = append(cw.Requests, &RequestContext{
+		Host:    reqhf.Get("Host"),
+		Path:    reqhf.Path,
+		Method:  reqhf.Method,
+		Chunked: false,
+	})
+	host := strings.Split(reqhf.Get("Host"), ".")
 	if len(host) < 1 {
 		_, err := conn.Write([]byte("HTTP/1.1 400 Bad Request\r\n\r\n"))
 		if err != nil {
@@ -166,7 +358,7 @@ func Handler(conn net.Conn) {
 	}
 
 	if slug == "ping" {
-		req, err := http.ReadRequest(reader)
+		req, err := http.ReadRequest(dstReader)
 		if err != nil {
 			log.Println("failed to parse HTTP request:", err)
 			return
@@ -218,39 +410,61 @@ func Handler(conn net.Conn) {
 		return
 	}
 
-	sshSession.HandleForwardedConnection(session.UserConnection{
-		Reader: reader,
-		Writer: conn,
-	}, sshSession.Conn)
+	forwardRequest(cw, reqhf, sshSession)
 	return
 }
 
-func peekUntilHeaders(reader *bufio.Reader, maxBytes int) ([]byte, error) {
-	var buf []byte
-	for {
-		n := len(buf) + 1
-		if n > maxBytes {
-			return buf, nil
-		}
-
-		peek, err := reader.Peek(n)
+func forwardRequest(cw *CustomWriter, initialRequest *RequestHeaderFactory, sshSession *session.SSHSession) {
+	cw.AddInteraction(sshSession.Interaction)
+	originHost, originPort := ParseAddr(cw.RemoteAddr.String())
+	payload := createForwardedTCPIPPayload(originHost, uint16(originPort), sshSession.Forwarder.GetForwardedPort())
+	channel, reqs, err := sshSession.Conn.OpenChannel("forwarded-tcpip", payload)
+	if err != nil {
+		log.Printf("Failed to open forwarded-tcpip channel: %v", err)
+		sendBadGatewayResponse(cw)
+		return
+	}
+	defer func(channel ssh.Channel) {
+		err := channel.Close()
 		if err != nil {
-			return nil, err
+			if errors.Is(err, io.EOF) {
+				sendBadGatewayResponse(cw)
+				return
+			}
+			log.Println("Failed to close connection:", err)
+			return
 		}
-		buf = peek
+	}(channel)
 
-		if bytes.Contains(buf, []byte("\r\n\r\n")) {
-			return buf, nil
+	go func() {
+		defer func() {
+			if r := recover(); r != nil {
+				log.Printf("Panic in request handler: %v", r)
+			}
+		}()
+		for req := range reqs {
+			err := req.Reply(false, nil)
+			if err != nil {
+				log.Printf("Failed to reply to request: %v", err)
+				return
+			}
 		}
+	}()
+
+	_, err = channel.Write(initialRequest.Finalize())
+	if err != nil {
+		log.Printf("Failed to write forwarded-tcpip:", err)
+		return
 	}
+
+	sshSession.HandleForwardedConnection(cw, channel, cw.RemoteAddr)
+	return
 }
 
-func parseHostFromHeader(data []byte) string {
-	lines := strings.Split(string(data), "\r\n")
-	for _, line := range lines {
-		if strings.HasPrefix(strings.ToLower(line), "host:") {
-			return strings.TrimSpace(strings.TrimPrefix(line, "Host:"))
-		}
+func sendBadGatewayResponse(writer io.Writer) {
+	_, err := writer.Write(BAD_GATEWAY_RESPONSE)
+	if err != nil {
+		log.Printf("failed to write Bad Gateway response: %v", err)
+		return
 	}
-	return ""
 }