LinkPort/internal/outbound/socks5.go

package outbound

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"linkport/internal/model"
	"net"
	"strconv"
	"time"
)

func dialSOCKS5(ctx context.Context, address string, o model.Outbound, keepAlive time.Duration) (_ net.Conn, resultErr error) {
	defer func() {
		if resultErr != nil {
			resultErr = fmt.Errorf("SOCKS5 negotiation failed: %w", resultErr)
		}
	}()
	d := net.Dialer{KeepAlive: keepAlive}
	done := Begin(ctx, "socks5_connect")
	c, err := d.DialContext(ctx, "tcp", net.JoinHostPort(o.Host, strconv.Itoa(o.Port)))
	done(err)
	if err != nil {
		return nil, fmt.Errorf("SOCKS5 connection failed: %w", err)
	}
	done = func(error) {}
	defer func() { done(resultErr) }()
	good := false
	defer func() {
		if !good {
			c.Close()
		}
	}()
	stop := context.AfterFunc(ctx, func() { c.Close() })
	defer stop()
	deadline, _ := ctx.Deadline()
	_ = c.SetDeadline(deadline)
	if err = negotiateSOCKS5(ctx, c, o); err != nil {
		return nil, err
	}
	done = Begin(ctx, "target_connect")
	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return nil, err
	}
	n, err := strconv.Atoi(port)
	if err != nil || n < 1 || n > 65535 {
		return nil, errors.New("invalid target port")
	}
	req := []byte{5, 1, 0}
	if ip := net.ParseIP(host); ip != nil {
		if v4 := ip.To4(); v4 != nil {
			req = append(req, 1)
			req = append(req, v4...)
		} else {
			req = append(req, 4)
			req = append(req, ip.To16()...)
		}
	} else {
		if len(host) < 1 || len(host) > 255 {
			return nil, errors.New("invalid target hostname")
		}
		req = append(req, 3, byte(len(host)))
		req = append(req, host...)
	}
	req = binary.BigEndian.AppendUint16(req, uint16(n))
	if err = write(c, req); err != nil {
		return nil, err
	}
	header := make([]byte, 4)
	if _, err = io.ReadFull(c, header); err != nil {
		return nil, err
	}
	if header[0] != 5 || header[1] != 0 || header[2] != 0 {
		return nil, fmt.Errorf("SOCKS5 target connection rejected (code %d)", header[1])
	}
	size := 0
	switch header[3] {
	case 1:
		size = 4
	case 4:
		size = 16
	case 3:
		var length [1]byte
		if _, err = io.ReadFull(c, length[:]); err != nil {
			return nil, err
		}
		size = int(length[0])
	default:
		return nil, errors.New("invalid SOCKS5 address")
	}
	if _, err = io.CopyN(io.Discard, c, int64(size+2)); err != nil {
		return nil, err
	}
	if !stop() || ctx.Err() != nil {
		return nil, context.Canceled
	}
	_ = c.SetDeadline(time.Time{})
	good = true
	return c, nil
}

func negotiateSOCKS5(ctx context.Context, connection net.Conn, profile model.Outbound) (resultErr error) {
	done := Begin(ctx, "socks5_negotiation")
	defer func() { done(resultErr) }()
	method := byte(0)
	if profile.Auth == "password" {
		method = 2
		if len(profile.Username) < 1 || len(profile.Username) > 255 || len(profile.Password) < 1 || len(profile.Password) > 255 {
			return errors.New("invalid SOCKS5 credentials")
		}
	}
	if err := write(connection, []byte{5, 1, method}); err != nil {
		return err
	}
	reply := make([]byte, 2)
	if _, err := io.ReadFull(connection, reply); err != nil {
		return err
	}
	if reply[0] != 5 || reply[1] != method {
		return model.AtConnectionStage("protocol_rejected", errors.New("SOCKS5 authentication method rejected"))
	}
	done(nil)
	if method == 2 {
		done = Begin(ctx, "socks5_authentication")
		req := append([]byte{1, byte(len(profile.Username))}, []byte(profile.Username)...)
		req = append(req, byte(len(profile.Password)))
		req = append(req, []byte(profile.Password)...)
		if err := write(connection, req); err != nil {
			return err
		}
		if _, err := io.ReadFull(connection, reply); err != nil {
			return err
		}
		if reply[0] != 1 || reply[1] != 0 {
			return model.AtConnectionStage("socks5_authentication", errors.New("SOCKS5 authentication failed"))
		}
		done(nil)
	}
	return nil
}

func write(w io.Writer, b []byte) error {
	for len(b) > 0 {
		n, err := w.Write(b)
		if err != nil {
			return err
		}
		if n == 0 {
			return io.ErrShortWrite
		}
		b = b[n:]
	}
	return nil
}