fix proxy protocol

service/service.go

@@ -0,0 +1,169 @@
+package service
+
+import (
+	"context"
+	"hash/crc32"
+	"net"
+	"time"
+
+	"github.com/go-gost/core/admission"
+	"github.com/go-gost/core/handler"
+	"github.com/go-gost/core/listener"
+	"github.com/go-gost/core/logger"
+	"github.com/go-gost/core/metrics"
+	"github.com/go-gost/core/recorder"
+	"github.com/go-gost/core/service"
+	sx "github.com/go-gost/x/internal/util/selector"
+	xmetrics "github.com/go-gost/x/metrics"
+)
+
+type options struct {
+	admission admission.Admission
+	recorders []recorder.RecorderObject
+	logger    logger.Logger
+}
+
+type Option func(opts *options)
+
+func AdmissionOption(admission admission.Admission) Option {
+	return func(opts *options) {
+		opts.admission = admission
+	}
+}
+
+func RecordersOption(recorders ...recorder.RecorderObject) Option {
+	return func(opts *options) {
+		opts.recorders = recorders
+	}
+}
+
+func LoggerOption(logger logger.Logger) Option {
+	return func(opts *options) {
+		opts.logger = logger
+	}
+}
+
+type defaultService struct {
+	name     string
+	listener listener.Listener
+	handler  handler.Handler
+	options  options
+}
+
+func NewService(name string, ln listener.Listener, h handler.Handler, opts ...Option) service.Service {
+	var options options
+	for _, opt := range opts {
+		opt(&options)
+	}
+	return &defaultService{
+		name:     name,
+		listener: ln,
+		handler:  h,
+		options:  options,
+	}
+}
+
+func (s *defaultService) Addr() net.Addr {
+	return s.listener.Addr()
+}
+
+func (s *defaultService) Close() error {
+	return s.listener.Close()
+}
+
+func (s *defaultService) Serve() error {
+	if v := xmetrics.GetGauge(
+		xmetrics.MetricServicesGauge,
+		metrics.Labels{}); v != nil {
+		v.Inc()
+		defer v.Dec()
+	}
+
+	var tempDelay time.Duration
+	for {
+		conn, e := s.listener.Accept()
+		if e != nil {
+			// TODO: remove Temporary checking
+			if ne, ok := e.(net.Error); ok && ne.Temporary() {
+				if tempDelay == 0 {
+					tempDelay = 1 * time.Second
+				} else {
+					tempDelay *= 2
+				}
+				if max := 5 * time.Second; tempDelay > max {
+					tempDelay = max
+				}
+				s.options.logger.Warnf("accept: %v, retrying in %v", e, tempDelay)
+				time.Sleep(tempDelay)
+				continue
+			}
+			s.options.logger.Errorf("accept: %v", e)
+			return e
+		}
+		tempDelay = 0
+
+		for _, rec := range s.options.recorders {
+			host := conn.RemoteAddr().String()
+			if h, _, _ := net.SplitHostPort(host); h != "" {
+				host = h
+			}
+			if rec.Record == recorder.RecorderServiceClientAddress {
+				if err := rec.Recorder.Record(context.Background(), []byte(host)); err != nil {
+					s.options.logger.Errorf("record %s: %v", rec.Record, err)
+				}
+			}
+		}
+		if s.options.admission != nil &&
+			!s.options.admission.Admit(conn.RemoteAddr().String()) {
+			conn.Close()
+			s.options.logger.Debugf("admission: %s is denied", conn.RemoteAddr())
+			continue
+		}
+
+		go func() {
+			if v := xmetrics.GetCounter(xmetrics.MetricServiceRequestsCounter,
+				metrics.Labels{"service": s.name}); v != nil {
+				v.Inc()
+			}
+
+			if v := xmetrics.GetGauge(xmetrics.MetricServiceRequestsInFlightGauge,
+				metrics.Labels{"service": s.name}); v != nil {
+				v.Inc()
+				defer v.Dec()
+			}
+
+			start := time.Now()
+			if v := xmetrics.GetObserver(xmetrics.MetricServiceRequestsDurationObserver,
+				metrics.Labels{"service": s.name}); v != nil {
+				defer func() {
+					v.Observe(float64(time.Since(start).Seconds()))
+				}()
+			}
+
+			ctx := sx.ContextWithHash(context.Background(), &sx.Hash{
+				Source: conn.RemoteAddr().String(),
+				Value:  ipHash(conn.RemoteAddr()),
+			})
+
+			if err := s.handler.Handle(ctx, conn); err != nil {
+				s.options.logger.Error(err)
+				if v := xmetrics.GetCounter(xmetrics.MetricServiceHandlerErrorsCounter,
+					metrics.Labels{"service": s.name}); v != nil {
+					v.Inc()
+				}
+			}
+		}()
+	}
+}
+
+func ipHash(addr net.Addr) uint64 {
+	if addr == nil {
+		return 0
+	}
+
+	host, _, _ := net.SplitHostPort(addr.String())
+	if ip := net.ParseIP(host); ip != nil {
+		return uint64(crc32.ChecksumIEEE(ip.To16()))
+	}
+	return 0
+}