listener.go 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899
  1. // Copyright 2017 fatedier, fatedier@gmail.com
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package net
  15. import (
  16. "fmt"
  17. "net"
  18. "sync"
  19. "github.com/fatedier/frp/utils/errors"
  20. "github.com/fatedier/frp/utils/log"
  21. )
  22. type Listener interface {
  23. Accept() (Conn, error)
  24. Close() error
  25. log.Logger
  26. }
  27. type LogListener struct {
  28. l net.Listener
  29. net.Listener
  30. log.Logger
  31. }
  32. func WrapLogListener(l net.Listener) Listener {
  33. return &LogListener{
  34. l: l,
  35. Listener: l,
  36. Logger: log.NewPrefixLogger(""),
  37. }
  38. }
  39. func (logL *LogListener) Accept() (Conn, error) {
  40. c, err := logL.l.Accept()
  41. return WrapConn(c), err
  42. }
  43. // Custom listener
  44. type CustomListener struct {
  45. conns chan Conn
  46. closed bool
  47. mu sync.Mutex
  48. log.Logger
  49. }
  50. func NewCustomListener() *CustomListener {
  51. return &CustomListener{
  52. conns: make(chan Conn, 64),
  53. Logger: log.NewPrefixLogger(""),
  54. }
  55. }
  56. func (l *CustomListener) Accept() (Conn, error) {
  57. conn, ok := <-l.conns
  58. if !ok {
  59. return nil, fmt.Errorf("listener closed")
  60. }
  61. conn.AddLogPrefix(l.GetPrefixStr())
  62. return conn, nil
  63. }
  64. func (l *CustomListener) PutConn(conn Conn) error {
  65. err := errors.PanicToError(func() {
  66. select {
  67. case l.conns <- conn:
  68. default:
  69. conn.Close()
  70. }
  71. })
  72. return err
  73. }
  74. func (l *CustomListener) Close() error {
  75. l.mu.Lock()
  76. defer l.mu.Unlock()
  77. if !l.closed {
  78. close(l.conns)
  79. l.closed = true
  80. }
  81. return nil
  82. }
  83. func (l *CustomListener) Addr() net.Addr {
  84. return (*net.TCPAddr)(nil)
  85. }