github.com/polarismesh/polaris@v1.17.8/common/conn/hook/listener.go (about)

     1  /**
     2   * Tencent is pleased to support the open source community by making Polaris available.
     3   *
     4   * Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
     5   *
     6   * Licensed under the BSD 3-Clause License (the "License");
     7   * you may not use this file except in compliance with the License.
     8   * You may obtain a copy of the License at
     9   *
    10   * https://opensource.org/licenses/BSD-3-Clause
    11   *
    12   * Unless required by applicable law or agreed to in writing, software distributed
    13   * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
    14   * CONDITIONS OF ANY KIND, either express or implied. See the License for the
    15   * specific language governing permissions and limitations under the License.
    16   */
    17  
    18  package connhook
    19  
    20  import (
    21  	"net"
    22  )
    23  
    24  // Hook Hook
    25  type Hook interface {
    26  	// OnAccept call when net.Conn accept
    27  	OnAccept(conn net.Conn)
    28  	// OnRelease call when net.Conn release
    29  	OnRelease(conn net.Conn)
    30  	// OnClose call when net.Listener close
    31  	OnClose()
    32  }
    33  
    34  func NewHookListener(l net.Listener, hooks ...Hook) net.Listener {
    35  	return &HookListener{
    36  		hooks:  hooks,
    37  		target: l,
    38  	}
    39  }
    40  
    41  // HookListener net.Listener can add hook
    42  type HookListener struct {
    43  	hooks  []Hook
    44  	target net.Listener
    45  }
    46  
    47  // Accept waits for and returns the next connection to the listener.
    48  func (l *HookListener) Accept() (net.Conn, error) {
    49  	conn, err := l.target.Accept()
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  
    54  	for i := range l.hooks {
    55  		l.hooks[i].OnAccept(conn)
    56  	}
    57  
    58  	return &Conn{Conn: conn, listener: l}, nil
    59  }
    60  
    61  // Close closes the listener.
    62  // Any blocked Accept operations will be unblocked and return errors.
    63  func (l *HookListener) Close() error {
    64  	for i := range l.hooks {
    65  		l.hooks[i].OnClose()
    66  	}
    67  	return l.target.Close()
    68  }
    69  
    70  // Addr returns the listener's network address.
    71  func (l *HookListener) Addr() net.Addr {
    72  	return l.target.Addr()
    73  }