Add middlewares for cozy filtering and logging

handle.go

@@ -86,16 +86,25 @@ func (sf *Server) handleRequest(write io.Writer, req *Request) error {
 		if sf.userConnectHandle != nil {
 			return sf.userConnectHandle(ctx, write, req)
 		}
+		if len(sf.userConnectMiddlewares) != 0 {
+			return sf.userConnectMiddlewares.Execute(ctx, write, req, sf.handleConnect)
+		}
 		return sf.handleConnect(ctx, write, req)
 	case statute.CommandBind:
 		if sf.userBindHandle != nil {
 			return sf.userBindHandle(ctx, write, req)
 		}
+		if len(sf.userBindMiddlewares) != 0 {
+			return sf.userBindMiddlewares.Execute(ctx, write, req, sf.handleBind)
+		}
 		return sf.handleBind(ctx, write, req)
 	case statute.CommandAssociate:
 		if sf.userAssociateHandle != nil {
 			return sf.userAssociateHandle(ctx, write, req)
 		}
+		if len(sf.userAssociateMiddlewares) != 0 {
+			return sf.userAssociateMiddlewares.Execute(ctx, write, req, sf.handleAssociate)
+		}
 		return sf.handleAssociate(ctx, write, req)
 	default:
 		if err := SendReply(write, statute.RepCommandNotSupported, nil); err != nil {

option.go

@@ -124,3 +124,46 @@ func WithAssociateHandle(h func(ctx context.Context, writer io.Writer, request *
 		s.userAssociateHandle = h
 	}
 }
+
+// Handler is used to handle a user's commands
+type Handler func(ctx context.Context, writer io.Writer, request *Request) error
+
+// WithMiddleware is used to add interceptors in chain
+type Middleware func(ctx context.Context, writer io.Writer, request *Request) error
+
+// MiddlewareChain is used to add interceptors in chain
+type MiddlewareChain []Middleware
+
+// Execute is used to add interceptors in chain
+func (m MiddlewareChain) Execute(ctx context.Context, writer io.Writer, request *Request, last Handler) error {
+	if len(m) == 0 {
+		return nil
+	}
+	for i := 0; i < len(m); i++ {
+		if err := m[i](ctx, writer, request); err != nil {
+			return err
+		}
+	}
+	return last(ctx, writer, request)
+}
+
+// WithConnectMiddleware is used to add interceptors in chain
+func WithConnectMiddleware(m Middleware) Option {
+	return func(s *Server) {
+		s.userConnectMiddlewares = append(s.userConnectMiddlewares, m)
+	}
+}
+
+// WithBindMiddleware is used to add interceptors in chain
+func WithBindMiddleware(m Middleware) Option {
+	return func(s *Server) {
+		s.userBindMiddlewares = append(s.userBindMiddlewares, m)
+	}
+}
+
+// WithAssociateMiddleware is used to add interceptors in chain
+func WithAssociateMiddleware(m Middleware) Option {
+	return func(s *Server) {
+		s.userAssociateMiddlewares = append(s.userAssociateMiddlewares, m)
+	}
+}

server.go

@@ -57,6 +57,10 @@ type Server struct {
 	userConnectHandle   func(ctx context.Context, writer io.Writer, request *Request) error
 	userBindHandle      func(ctx context.Context, writer io.Writer, request *Request) error
 	userAssociateHandle func(ctx context.Context, writer io.Writer, request *Request) error
+	// user's middleware
+	userConnectMiddlewares   MiddlewareChain
+	userBindMiddlewares      MiddlewareChain
+	userAssociateMiddlewares MiddlewareChain
 }
 
 // NewServer creates a new Server