default_proxy.go 1014 B

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. package cproxy
  2. import (
  3. "io"
  4. "sync"
  5. )
  6. type defaultProxy struct {
  7. client Socket
  8. server Socket
  9. waiter *sync.WaitGroup
  10. }
  11. func newProxy(client, server Socket) *defaultProxy {
  12. waiter := &sync.WaitGroup{}
  13. waiter.Add(2) // wait on both client->server and server->client streams
  14. return &defaultProxy{
  15. waiter: waiter,
  16. client: client,
  17. server: server,
  18. }
  19. }
  20. func (this *defaultProxy) Proxy() {
  21. go this.streamAndClose(this.client, this.server)
  22. go this.streamAndClose(this.server, this.client)
  23. this.closeSockets()
  24. }
  25. func (this *defaultProxy) streamAndClose(reader, writer Socket) {
  26. _, _ = io.Copy(writer, reader)
  27. tryCloseRead(reader)
  28. tryCloseWrite(writer)
  29. this.waiter.Done()
  30. }
  31. func tryCloseRead(socket Socket) {
  32. if tcp, ok := socket.(tcpSocket); ok {
  33. _ = tcp.CloseRead()
  34. }
  35. }
  36. func tryCloseWrite(socket Socket) {
  37. if tcp, ok := socket.(tcpSocket); ok {
  38. _ = tcp.CloseWrite()
  39. }
  40. }
  41. func (this *defaultProxy) closeSockets() {
  42. this.waiter.Wait()
  43. _ = this.client.Close()
  44. _ = this.server.Close()
  45. }