Merge pull request #71 from EvgenyOvsov/add_middlewares

Add middlewares for cozy filtering and logging

handle.go

@@ -80,29 +80,40 @@ func (sf *Server) handleRequest(write io.Writer, req *Request) error {
  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 @@ 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

@@ -58,6 +58,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