使用Golang通过chan转发DNS请求实现DNS代理

要实现通过通道转发DNS请求的代理，关键在于正确管理请求ID和响应通道的映射关系。以下是改进后的实现方案：

package main

import (
    "sync"
    "time"
    "github.com/google/gopacket"
    "github.com/google/gopacket/layers"
    "github.com/miekg/dns"
    "github.com/google/gopacket/pcap"
)

type DNSRequest struct {
    ID       uint16
    Response chan *layers.DNS
    Timeout  time.Time
}

type DNSProxy struct {
    requests     map[uint16]*DNSRequest
    requestsLock sync.RWMutex
    requestChan  chan *dns.Msg
    responseChan chan *layers.DNS
}

func NewDNSProxy() *DNSProxy {
    return &DNSProxy{
        requests:    make(map[uint16]*DNSRequest),
        requestChan: make(chan *dns.Msg, 100),
        responseChan: make(chan *layers.DNS, 100),
    }
}

func (p *DNSProxy) handleRequest(w dns.ResponseWriter, req *dns.Msg) {
    // 将请求放入通道，异步处理
    p.requestChan <- req
    
    // 创建响应通道
    respChan := make(chan *layers.DNS, 1)
    
    p.requestsLock.Lock()
    p.requests[req.Id] = &DNSRequest{
        ID:       req.Id,
        Response: respChan,
        Timeout:  time.Now().Add(5 * time.Second),
    }
    p.requestsLock.Unlock()
    
    // 等待响应或超时
    select {
    case dnsResp := <-respChan:
        // 将layers.DNS转换为dns.Msg
        resp := p.convertToDNSMsg(dnsResp)
        w.WriteMsg(resp)
    case <-time.After(5 * time.Second):
        // 返回超时响应
        resp := new(dns.Msg)
        resp.SetRcode(req, dns.RcodeServerFailure)
        w.WriteMsg(resp)
    }
    
    // 清理
    p.requestsLock.Lock()
    delete(p.requests, req.Id)
    p.requestsLock.Unlock()
}

func (p *DNSProxy) StartRequestProcessor() {
    go func() {
        for req := range p.requestChan {
            go p.sendDNSRequest(req)
        }
    }()
}

func (p *DNSProxy) sendDNSRequest(req *dns.Msg) {
    // 打开WiFi接口发送请求
    handle, err := pcap.OpenLive(`\Device\NPF_{<Device-Hardw-Id>}`, 2048, false, 0)
    if err != nil {
        return
    }
    defer handle.Close()
    
    // 构建DNS请求包
    dnsLayer := &layers.DNS{
        ID:           req.Id,
        QR:           false,
        OpCode:       layers.DNSOpCodeQuery,
        RD:           true,
        QDCount:      1,
        Questions: []layers.DNSQuestion{
            {
                Name:  []byte(req.Question[0].Name),
                Type:  layers.DNSTypeA,
                Class: layers.DNSClassIN,
            },
        },
    }
    
    // 构建完整的数据包
    buffer := gopacket.NewSerializeBuffer()
    options := gopacket.SerializeOptions{
        ComputeChecksums: true,
        FixLengths:       true,
    }
    
    // 构建网络层
    eth := &layers.Ethernet{
        SrcMAC:       net.HardwareAddr{0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
        DstMAC:       net.HardwareAddr{0xff, 0xff, 0xff, 0xff, 0xff, 0xff},
        EthernetType: layers.EthernetTypeIPv4,
    }
    
    ip := &layers.IPv4{
        Version:  4,
        TTL:      64,
        SrcIP:    net.IP{192, 168, 1, 100},
        DstIP:    net.IP{8, 8, 8, 8},
        Protocol: layers.IPProtocolUDP,
    }
    
    udp := &layers.UDP{
        SrcPort: 5353,
        DstPort: 53,
    }
    udp.SetNetworkLayerForChecksum(ip)
    
    gopacket.SerializeLayers(buffer, options,
        eth,
        ip,
        udp,
        dnsLayer,
    )
    
    handle.WritePacketData(buffer.Bytes())
}

func (p *DNSProxy) ListenToWifi() {
    handle, err := pcap.OpenLive(`\Device\NPF_{<Device-Hardw-Id>}`, 2048, false, 100*time.Millisecond)
    if err != nil {
        log.Fatal(err)
    }
    defer handle.Close()
    
    packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
    for packet := range packetSource.Packets() {
        dnsLayer := packet.Layer(layers.LayerTypeDNS)
        if dnsLayer == nil {
            continue
        }
        
        dns, ok := dnsLayer.(*layers.DNS)
        if !ok || !dns.QR {
            continue
        }
        
        // 将响应发送到通道
        p.responseChan <- dns
    }
}

func (p *DNSProxy) StartResponseProcessor() {
    go func() {
        for dnsResp := range p.responseChan {
            p.requestsLock.RLock()
            if req, exists := p.requests[dnsResp.ID]; exists {
                select {
                case req.Response <- dnsResp:
                    // 成功发送响应
                default:
                    // 通道已满，丢弃响应
                }
            }
            p.requestsLock.RUnlock()
        }
    }()
}

func (p *DNSProxy) convertToDNSMsg(dnsResp *layers.DNS) *dns.Msg {
    msg := new(dns.Msg)
    msg.Id = dnsResp.ID
    msg.Response = true
    
    // 转换问题部分
    for _, q := range dnsResp.Questions {
        msg.Question = append(msg.Question, dns.Question{
            Name:   string(q.Name),
            Qtype:  uint16(q.Type),
            Qclass: uint16(q.Class),
        })
    }
    
    // 转换答案部分
    for _, a := range dnsResp.Answers {
        rr := &dns.A{
            Hdr: dns.RR_Header{
                Name:   string(a.Name),
                Rrtype: uint16(a.Type),
                Class:  uint16(a.Class),
                Ttl:    a.TTL,
            },
            A: net.IP(a.IP),
        }
        msg.Answer = append(msg.Answer, rr)
    }
    
    return msg
}

func (p *DNSProxy) CleanupExpiredRequests() {
    go func() {
        ticker := time.NewTicker(1 * time.Second)
        defer ticker.Stop()
        
        for range ticker.C {
            p.requestsLock.Lock()
            now := time.Now()
            for id, req := range p.requests {
                if now.After(req.Timeout) {
                    close(req.Response)
                    delete(p.requests, id)
                }
            }
            p.requestsLock.Unlock()
        }
    }()
}

func main() {
    proxy := NewDNSProxy()
    
    // 启动DNS服务器
    server := &dns.Server{
        Addr:    ":53",
        Net:     "udp",
        Handler: dns.HandlerFunc(proxy.handleRequest),
    }
    
    go server.ListenAndServe()
    
    // 启动各个处理器
    proxy.StartRequestProcessor()
    proxy.StartResponseProcessor()
    proxy.CleanupExpiredRequests()
    
    // 监听WiFi接口
    go proxy.ListenToWifi()
    
    select {}
}

这个实现方案解决了你提到的问题：

1. 并发请求处理：使用sync.RWMutex保护请求映射，每个请求有独立的响应通道
2. 超时处理：每个请求设置5秒超时，避免阻塞
3. 异步处理：请求和响应通过通道异步传递，不会阻塞监听
4. 内存管理：定期清理过期的请求，防止内存泄漏

关键改进点：

• 使用请求ID作为键的映射表来跟踪请求
• 每个请求创建独立的响应通道
• 异步处理请求和响应
• 添加超时机制和清理机制