util.go 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  1. package util
  2. import (
  3. "crypto/rand"
  4. "encoding/hex"
  5. "fmt"
  6. "log"
  7. "net"
  8. "sync"
  9. "time"
  10. )
  11. func GenPSK(length int) ([]byte, error) {
  12. b := make([]byte, length)
  13. _, err := rand.Read(b)
  14. if err != nil {
  15. return nil, fmt.Errorf("random bytes generation failed: %w", err)
  16. }
  17. return b, nil
  18. }
  19. func GenPSKHex(length int) (string, error) {
  20. b, err := GenPSK(length)
  21. if err != nil {
  22. return "", fmt.Errorf("can't generate hex key: %w", err)
  23. }
  24. return hex.EncodeToString(b), nil
  25. }
  26. func PSKFromHex(input string) ([]byte, error) {
  27. return hex.DecodeString(input)
  28. }
  29. func isTimeout(err error) bool {
  30. if timeoutErr, ok := err.(interface {
  31. Timeout() bool
  32. }); ok {
  33. return timeoutErr.Timeout()
  34. }
  35. return false
  36. }
  37. func isTemporary(err error) bool {
  38. if timeoutErr, ok := err.(interface {
  39. Temporary() bool
  40. }); ok {
  41. return timeoutErr.Temporary()
  42. }
  43. return false
  44. }
  45. const (
  46. MaxPktBuf = 65536
  47. )
  48. func PairConn(left, right net.Conn, idleTimeout time.Duration, staleMode StaleMode) {
  49. var wg sync.WaitGroup
  50. tracker := newTracker(staleMode)
  51. copier := func(dst, src net.Conn, label bool) {
  52. defer wg.Done()
  53. defer dst.Close()
  54. buf := make([]byte, MaxPktBuf)
  55. for {
  56. if err := src.SetReadDeadline(time.Now().Add(idleTimeout)); err != nil {
  57. log.Printf("can't update deadline for connection: %v", err)
  58. break
  59. }
  60. n, err := src.Read(buf)
  61. if err != nil {
  62. if isTimeout(err) {
  63. // hit read deadline
  64. if tracker.handleTimeout(label) {
  65. // not stale conn
  66. continue
  67. } else {
  68. log.Printf("dropping stale connection %s <=> %s", src.LocalAddr(), src.RemoteAddr())
  69. }
  70. } else {
  71. // any other error
  72. if isTemporary(err) {
  73. log.Printf("ignoring temporary error during read from %s: %v", src.RemoteAddr(), err)
  74. continue
  75. }
  76. log.Printf("read from %s error: %v", src.RemoteAddr(), err)
  77. }
  78. break
  79. }
  80. tracker.notify(label)
  81. _, err = dst.Write(buf[:n])
  82. if err != nil {
  83. log.Printf("write to %s error: %v", dst.RemoteAddr(), err)
  84. break
  85. }
  86. }
  87. }
  88. wg.Add(2)
  89. go copier(left, right, false)
  90. go copier(right, left, true)
  91. wg.Wait()
  92. }