plugin.go 2.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192
  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 plugin
  15. import (
  16. "fmt"
  17. "io"
  18. "net"
  19. "sync"
  20. "github.com/fatedier/golib/errors"
  21. )
  22. // Creators is used for create plugins to handle connections.
  23. var creators = make(map[string]CreatorFn)
  24. // params has prefix "plugin_"
  25. type CreatorFn func(params map[string]string) (Plugin, error)
  26. func Register(name string, fn CreatorFn) {
  27. creators[name] = fn
  28. }
  29. func Create(name string, params map[string]string) (p Plugin, err error) {
  30. if fn, ok := creators[name]; ok {
  31. p, err = fn(params)
  32. } else {
  33. err = fmt.Errorf("plugin [%s] is not registered", name)
  34. }
  35. return
  36. }
  37. type Plugin interface {
  38. Name() string
  39. // extraBufToLocal will send to local connection first, then join conn with local connection
  40. Handle(conn io.ReadWriteCloser, realConn net.Conn, extraBufToLocal []byte)
  41. Close() error
  42. }
  43. type Listener struct {
  44. conns chan net.Conn
  45. closed bool
  46. mu sync.Mutex
  47. }
  48. func NewProxyListener() *Listener {
  49. return &Listener{
  50. conns: make(chan net.Conn, 64),
  51. }
  52. }
  53. func (l *Listener) Accept() (net.Conn, error) {
  54. conn, ok := <-l.conns
  55. if !ok {
  56. return nil, fmt.Errorf("listener closed")
  57. }
  58. return conn, nil
  59. }
  60. func (l *Listener) PutConn(conn net.Conn) error {
  61. err := errors.PanicToError(func() {
  62. l.conns <- conn
  63. })
  64. return err
  65. }
  66. func (l *Listener) Close() error {
  67. l.mu.Lock()
  68. defer l.mu.Unlock()
  69. if !l.closed {
  70. close(l.conns)
  71. l.closed = true
  72. }
  73. return nil
  74. }
  75. func (l *Listener) Addr() net.Addr {
  76. return (*net.TCPAddr)(nil)
  77. }