relay.go 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  1. package resolver
  2. import (
  3. "context"
  4. "encoding/binary"
  5. "io"
  6. "net"
  7. "time"
  8. "github.com/metacubex/mihomo/common/pool"
  9. D "github.com/miekg/dns"
  10. )
  11. const DefaultDnsReadTimeout = time.Second * 10
  12. const DefaultDnsRelayTimeout = time.Second * 5
  13. const SafeDnsPacketSize = 2 * 1024 // safe size which is 1232 from https://dnsflagday.net/2020/, so 2048 is enough
  14. func RelayDnsConn(ctx context.Context, conn net.Conn, readTimeout time.Duration) error {
  15. buff := pool.Get(pool.UDPBufferSize)
  16. defer func() {
  17. _ = pool.Put(buff)
  18. _ = conn.Close()
  19. }()
  20. for {
  21. if readTimeout > 0 {
  22. _ = conn.SetReadDeadline(time.Now().Add(readTimeout))
  23. }
  24. length := uint16(0)
  25. if err := binary.Read(conn, binary.BigEndian, &length); err != nil {
  26. break
  27. }
  28. if int(length) > len(buff) {
  29. break
  30. }
  31. n, err := io.ReadFull(conn, buff[:length])
  32. if err != nil {
  33. break
  34. }
  35. err = func() error {
  36. ctx, cancel := context.WithTimeout(ctx, DefaultDnsRelayTimeout)
  37. defer cancel()
  38. inData := buff[:n]
  39. msg, err := relayDnsPacket(ctx, inData, buff, 0)
  40. if err != nil {
  41. return err
  42. }
  43. err = binary.Write(conn, binary.BigEndian, uint16(len(msg)))
  44. if err != nil {
  45. return err
  46. }
  47. _, err = conn.Write(msg)
  48. if err != nil {
  49. return err
  50. }
  51. return nil
  52. }()
  53. if err != nil {
  54. return err
  55. }
  56. }
  57. return nil
  58. }
  59. func relayDnsPacket(ctx context.Context, payload []byte, target []byte, maxSize int) ([]byte, error) {
  60. msg := &D.Msg{}
  61. if err := msg.Unpack(payload); err != nil {
  62. return nil, err
  63. }
  64. r, err := ServeMsg(ctx, msg)
  65. if err != nil {
  66. m := new(D.Msg)
  67. m.SetRcode(msg, D.RcodeServerFailure)
  68. return m.PackBuffer(target)
  69. }
  70. r.SetRcode(msg, r.Rcode)
  71. if maxSize > 0 {
  72. r.Truncate(maxSize)
  73. }
  74. r.Compress = true
  75. return r.PackBuffer(target)
  76. }
  77. // RelayDnsPacket will truncate udp message up to SafeDnsPacketSize
  78. func RelayDnsPacket(ctx context.Context, payload []byte, target []byte) ([]byte, error) {
  79. return relayDnsPacket(ctx, payload, target, SafeDnsPacketSize)
  80. }