feat: add header factory

server/header.go

@@ -0,0 +1,163 @@
+package server
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"io"
+	"strings"
+)
+
+type HeaderManager interface {
+	Get(key string) []byte
+	Set(key string, value []byte)
+	Remove(key string)
+	Finalize() []byte
+}
+
+type ResponseHeaderFactory struct {
+	startLine []byte
+	headers   map[string]string
+}
+
+type RequestHeaderFactory struct {
+	Method    string
+	Path      string
+	Version   string
+	startLine []byte
+	headers   map[string]string
+}
+
+func NewRequestHeaderFactory(r io.Reader) (*RequestHeaderFactory, error) {
+	br := bufio.NewReader(r)
+	header := &RequestHeaderFactory{
+		headers: make(map[string]string),
+	}
+
+	startLine, err := br.ReadString('\n')
+	if err != nil {
+		return nil, err
+	}
+	startLine = strings.TrimRight(startLine, "\r\n")
+	header.startLine = []byte(startLine)
+
+	parts := strings.Split(startLine, " ")
+	if len(parts) < 3 {
+		return nil, fmt.Errorf("invalid request line")
+	}
+
+	header.Method = parts[0]
+	header.Path = parts[1]
+	header.Version = parts[2]
+
+	for {
+		line, err := br.ReadString('\n')
+		if err != nil {
+			return nil, err
+		}
+		line = strings.TrimRight(line, "\r\n")
+
+		if line == "" {
+			break
+		}
+
+		kv := strings.SplitN(line, ":", 2)
+		if len(kv) != 2 {
+			continue
+		}
+		header.headers[strings.TrimSpace(kv[0])] = strings.TrimSpace(kv[1])
+	}
+
+	return header, nil
+}
+
+func NewResponseHeaderFactory(startLine []byte) *ResponseHeaderFactory {
+	header := &ResponseHeaderFactory{
+		startLine: nil,
+		headers:   make(map[string]string),
+	}
+	lines := bytes.Split(startLine, []byte("\r\n"))
+	if len(lines) == 0 {
+		return header
+	}
+	header.startLine = lines[0]
+	for _, h := range lines[1:] {
+		if len(h) == 0 {
+			continue
+		}
+
+		parts := bytes.SplitN(h, []byte(":"), 2)
+		if len(parts) < 2 {
+			continue
+		}
+
+		key := parts[0]
+		val := bytes.TrimSpace(parts[1])
+		header.headers[string(key)] = string(val)
+	}
+	return header
+}
+
+func (resp *ResponseHeaderFactory) Get(key string) string {
+	return resp.headers[key]
+}
+
+func (resp *ResponseHeaderFactory) Set(key string, value string) {
+	resp.headers[key] = value
+}
+
+func (resp *ResponseHeaderFactory) Remove(key string) {
+	delete(resp.headers, key)
+}
+
+func (resp *ResponseHeaderFactory) Finalize() []byte {
+	var buf bytes.Buffer
+
+	buf.Write(resp.startLine)
+	buf.WriteString("\r\n")
+
+	for key, val := range resp.headers {
+		buf.WriteString(key)
+		buf.WriteString(": ")
+		buf.WriteString(val)
+		buf.WriteString("\r\n")
+	}
+
+	buf.WriteString("\r\n")
+	return buf.Bytes()
+}
+
+func (req *RequestHeaderFactory) Get(key string) string {
+	val, ok := req.headers[key]
+	if !ok {
+		return ""
+	}
+	return val
+}
+
+func (req *RequestHeaderFactory) Set(key string, value string) {
+	req.headers[key] = value
+}
+
+func (req *RequestHeaderFactory) Remove(key string) {
+	delete(req.headers, key)
+}
+
+func (req *RequestHeaderFactory) Finalize() []byte {
+	var buf bytes.Buffer
+
+	buf.Write(req.startLine)
+	buf.WriteString("\r\n")
+
+	req.headers["X-HF"] = "modified"
+
+	for key, val := range req.headers {
+		buf.WriteString(key)
+		buf.WriteString(": ")
+		buf.WriteString(val)
+		buf.WriteString("\r\n")
+	}
+
+	buf.WriteString("\r\n")
+	return buf.Bytes()
+}

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 ""
 }

server/https.go

@@ -45,14 +45,31 @@ func NewHTTPSServer() error {
 }
 
 func HandlerTLS(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 {
@@ -70,7 +87,7 @@ func HandlerTLS(conn net.Conn) {
 	slug := host[0]
 
 	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
@@ -121,10 +138,6 @@ func HandlerTLS(conn net.Conn) {
 		}
 		return
 	}
-
-	sshSession.HandleForwardedConnection(session.UserConnection{
-		Reader: reader,
-		Writer: conn,
-	}, sshSession.Conn)
+	forwardRequest(cw, reqhf, sshSession)
 	return
 }

server/server.go

@@ -1,12 +1,16 @@
 package server
 
 import (
+	"bytes"
+	"encoding/binary"
 	"fmt"
-	"golang.org/x/crypto/ssh"
 	"log"
 	"net"
 	"net/http"
+	"strconv"
 	"tunnel_pls/utils"
+
+	"golang.org/x/crypto/ssh"
 )
 
 type Server struct {
@@ -54,3 +58,41 @@ func (s *Server) Start() {
 		go s.handleConnection(conn)
 	}
 }
+
+func createForwardedTCPIPPayload(host string, originPort, port uint16) []byte {
+	var buf bytes.Buffer
+
+	writeSSHString(&buf, "localhost")
+	err := binary.Write(&buf, binary.BigEndian, uint32(port))
+	if err != nil {
+		log.Printf("Failed to write string to buffer: %v", err)
+		return nil
+	}
+	writeSSHString(&buf, host)
+	err = binary.Write(&buf, binary.BigEndian, uint32(originPort))
+	if err != nil {
+		log.Printf("Failed to write string to buffer: %v", err)
+		return nil
+	}
+
+	return buf.Bytes()
+}
+
+func writeSSHString(buffer *bytes.Buffer, str string) {
+	err := binary.Write(buffer, binary.BigEndian, uint32(len(str)))
+	if err != nil {
+		log.Printf("Failed to write string to buffer: %v", err)
+		return
+	}
+	buffer.WriteString(str)
+}
+
+func ParseAddr(addr string) (string, uint32) {
+	host, portStr, err := net.SplitHostPort(addr)
+	if err != nil {
+		log.Printf("Failed to parse origin address: %s from address %s", err.Error(), addr)
+		return "0.0.0.0", uint32(0)
+	}
+	port, _ := strconv.Atoi(portStr)
+	return host, uint32(port)
+}