things-go / go-socks5

socks5 server in pure Golang with much custom optional. Full TCP/UDP and IPv4/IPv6 support.
MIT License
412 stars 72 forks source link

[FEATURE REQUEST] How to Implement Connection Count Statistics ? #36

Closed gopkg-dev closed 1 year ago

gopkg-dev commented 1 year ago

Description: I have added functionality to track active connections in the handleConnect method of the go-socks5 package. I would like to request a code review and suggestions for improving the implementation.

https://github.com/things-go/go-socks5/blob/ea3e2a06cb41c7fa992aff40673a0a8211effa3f/handle.go#L108-L150


type Server struct {
    // ...
    ActiveConnections int32 // Number of active connections
    // ...
}

// GetActiveConnections returns the current count of active connections
func (sf *Server) GetActiveConnections() int32 {
    return atomic.LoadInt32(&sf.ActiveConnections)
}

// handleConnect is used to handle a connect command
func (sf *Server) handleConnect(ctx context.Context, writer io.Writer, request *Request) error {
    // Attempt to connect
    dial := sf.dial
    if dial == nil {
        dial = func(ctx context.Context, net_, addr string) (net.Conn, error) {
            return net.Dial(net_, addr)
        }
    }
    target, err := dial(ctx, "tcp", request.DestAddr.String())
    if err != nil {
        msg := err.Error()
        resp := statute.RepHostUnreachable
        if strings.Contains(msg, "refused") {
            resp = statute.RepConnectionRefused
        } else if strings.Contains(msg, "network is unreachable") {
            resp = statute.RepNetworkUnreachable
        }
        if err := SendReply(writer, resp, nil); err != nil {
            return fmt.Errorf("failed to send reply, %v", err)
        }
        return fmt.Errorf("connect to %v failed, %v", request.RawDestAddr, err)
    }

    // Increment the connection count when establishing a connection
    atomic.AddInt32(&sf.ActiveConnections, 1)
    defer func() {
        target.Close() //nolint: errcheck
        atomic.AddInt32(&sf.ActiveConnections, -1)
    }()

    // Send success
    if err := SendReply(writer, statute.RepSuccess, target.LocalAddr()); err != nil {
        return fmt.Errorf("failed to send reply, %v", err)
    }

    // Start proxying
    errCh := make(chan error, 2)
    sf.goFunc(func() { errCh <- sf.Proxy(target, request.Reader) })
    sf.goFunc(func() { errCh <- sf.Proxy(writer, target) })
    // Wait
    for i := 0; i < 2; i++ {
        e := <-errCh
        if e != nil {
            // return from this function closes target (and conn).
            return e
        }
    }
    return nil
}
thinkgos commented 1 year ago

@gopkg-dev suggest add prestart and poststop function in handleRequest,. custom for user defined it's owned feature.