Add middlewares for cozy filtering and logging

handle.go

@@ -80,29 +80,40 @@
 		return fmt.Errorf("bind to %v blocked by rules", req.RawDestAddr)
 	}
 
+	var last Handler
 	// Switch on the command
 	switch req.Command {
 	case statute.CommandConnect:
+		last = sf.handleConnect
 		if sf.userConnectHandle != nil {
-			return sf.userConnectHandle(ctx, write, req)
+			last = sf.userConnectHandle
+		}
+		if len(sf.userConnectMiddlewares) != 0 {
+			return sf.userConnectMiddlewares.Execute(ctx, write, req, last)
 		}
-		return sf.handleConnect(ctx, write, req)
 	case statute.CommandBind:
+		last = sf.handleBind
 		if sf.userBindHandle != nil {
-			return sf.userBindHandle(ctx, write, req)
+			last = sf.userBindHandle
+		}
+		if len(sf.userBindMiddlewares) != 0 {
+			return sf.userBindMiddlewares.Execute(ctx, write, req, last)
 		}
-		return sf.handleBind(ctx, write, req)
 	case statute.CommandAssociate:
+		last = sf.handleAssociate
 		if sf.userAssociateHandle != nil {
-			return sf.userAssociateHandle(ctx, write, req)
+			last = sf.userAssociateHandle
+		}
+		if len(sf.userAssociateMiddlewares) != 0 {
+			return sf.userAssociateMiddlewares.Execute(ctx, write, req, last)
 		}
-		return sf.handleAssociate(ctx, write, req)
 	default:
 		if err := SendReply(write, statute.RepCommandNotSupported, nil); err != nil {
 			return fmt.Errorf("failed to send reply, %v", err)
 		}
 		return fmt.Errorf("unsupported command[%v]", req.Command)
 	}
+	return last(ctx, write, req)
 }
 
 // handleConnect is used to handle a connect command

option.go

@@ -124,3 +124,46 @@
 		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