v0.5.0 (#3)

authentication.go

@@ -39,6 +39,8 @@ func (s *Server) usernamePasswordAuthenticate(ctx context.Context, writer io.Wri
 		return
 	}
 
+	ctx = contextWithUsername(ctx, string(username))
+
 	passwordLen, err := reader.ReadByte()
 	if err != nil {
 		s.logger.Error(ctx, "failed to read password length: "+err.Error())

context.go

@@ -4,7 +4,10 @@ import "context"
 
 type ctxKey int
 
-const remoteAddressKey ctxKey = iota
+const (
+	remoteAddressKey ctxKey = iota
+	usernameKey
+)
 
 func contextWithRemoteAddress(ctx context.Context, address string) context.Context {
 	return context.WithValue(ctx, remoteAddressKey, address)
@@ -14,3 +17,12 @@ func RemoteAddressFromContext(ctx context.Context) (string, bool) {
 	value, ok := ctx.Value(remoteAddressKey).(string)
 	return value, ok
 }
+
+func contextWithUsername(ctx context.Context, username string) context.Context {
+	return context.WithValue(ctx, usernameKey, username)
+}
+
+func UsernameFromContext(ctx context.Context, username string) (string, bool) {
+	value, ok := ctx.Value(usernameKey).(string)
+	return value, ok
+}

error.go

@@ -24,6 +24,6 @@ func connectionRefusedError(err error) bool {
 	return strings.Contains(err.Error(), "refused")
 }
 
-func closedListenerError(e error) bool {
-	return errors.Is(e, net.ErrClosed)
+func closedListenerError(err error) bool {
+	return errors.Is(err, net.ErrClosed)
 }

metrics.go

@@ -0,0 +1,13 @@
+package socks5
+
+import "context"
+
+type Metrics interface {
+	UploadBytes(ctx context.Context, n int64)
+	DownloadBytes(ctx context.Context, n int64)
+}
+
+type nopMetrics struct{}
+
+func (m *nopMetrics) UploadBytes(_ context.Context, _ int64)   {}
+func (m *nopMetrics) DownloadBytes(_ context.Context, _ int64) {}

options.go

@@ -17,6 +17,7 @@ type Options struct {
 	Logger             Logger            // default: stdoutLogger
 	Store              Store             // default: defaultStore
 	Driver             Driver            // default: defaultDriver
+	Metrics            Metrics           // default: nopMetrics
 }
 
 func (o Options) authMethods() map[byte]struct{} {
@@ -66,5 +67,9 @@ func optsWithDefaults(opts *Options) *Options {
 		}
 	}
 
+	if opts.Metrics == nil {
+		opts.Metrics = &nopMetrics{}
+	}
+
 	return opts
 }

relay.go

@@ -6,13 +6,13 @@ type closeWriter interface {
 	CloseWrite() error
 }
 
-func relay(dst io.Writer, src io.Reader) error {
-	_, err := io.Copy(dst, src)
+func relay(dst io.Writer, src io.Reader) (int64, error) {
+	n, err := io.Copy(dst, src)
 
 	if writer, ok := dst.(closeWriter); ok {
 		// send EOF for next io.Copy
 		writer.CloseWrite()
 	}
 
-	return err
+	return n, err
 }

server.go

@@ -19,6 +19,7 @@ type Server struct {
 	logger        Logger
 	store         Store
 	driver        Driver
+	metrics       Metrics
 	active        chan struct{}
 	done          chan struct{}
 	closeListener func() error
@@ -34,11 +35,12 @@ func New(opts *Options) *Server {
 			getPasswordTimeout: opts.GetPasswordTimeout,
 			authMethods:        opts.authMethods(),
 		},
-		logger: opts.Logger,
-		store:  opts.Store,
-		driver: opts.Driver,
-		active: make(chan struct{}),
-		done:   make(chan struct{}),
+		logger:  opts.Logger,
+		store:   opts.Store,
+		driver:  opts.Driver,
+		metrics: opts.Metrics,
+		active:  make(chan struct{}),
+		done:    make(chan struct{}),
 	}
 }
 

socks5.go

@@ -167,11 +167,15 @@ func (s *Server) connect(ctx context.Context, writer io.Writer, reader *bufio.Re
 	var g errgroup.Group
 
 	g.Go(func() error {
-		return relay(target, reader)
+		n, err := relay(target, reader)
+		s.metrics.UploadBytes(ctx, n)
+		return err
 	})
 
 	g.Go(func() error {
-		return relay(writer, target)
+		n, err := relay(writer, target)
+		s.metrics.DownloadBytes(ctx, n)
+		return err
 	})
 
 	if err = g.Wait(); err != nil {