tcp流量转发到不同的后端

jiangyd · 2022-09-21 18:01:22 · 2596 次点击

想把流量发个多个后端，每个后端都可以收到内容，但我写的，每次都只有一个后端收到内容，随机的

package main

import (
    "flag"
    "fmt"
    "io"
    "net"
    "strings"
)

var tcp_conn []net.Conn

func main() {
    var address *string
    var dst *string
    address = flag.String("listen_address", "127.0.0.1:8080", "listen address")
    dst = flag.String("target_address", "172.16.5.9:18061,172.16.5.9:18060", "目标地址,多个地址逗号分割")
    flag.Parse()
    listener, err := net.Listen("tcp", *address)
    if err != nil {
        fmt.Println(err)
    }
    ips := strings.Split(*dst, ",")
    if len(ips) == 0 {
        fmt.Println("目标地址不能为空")
        return
    }
    for _, ip := range ips {
        conn, err := net.Dial("tcp", ip)
        if err != nil {
            fmt.Println(err)
            continue
        }
        tcp_conn = append(tcp_conn, conn)
    }
    for {
        conn, err := listener.Accept()
        if err != nil {
            continue
        }
        go handler(conn)
    }

}

func handler(conn net.Conn) {
    // defer conn.Close()
    for _, dst_conn := range tcp_conn {
        go io.Copy(dst_conn, conn)
        go io.Copy(conn, dst_conn)
    }
}

该如何改呢，哪位大佬帮忙看看

阅读全文

查看全部 11 个评论

zzustu

凭感觉写的，未经优化未经测试，仅供参考

package main

import (
    "io"
    "net"
    "time"
)

func main() {
    backends := []string{"172.16.5.9:18060", "172.16.5.9:18061", "172.16.5.9:18062"}

    lis, err := net.Listen("tcp", ":8080")
    if err != nil {
        panic(err)
    }

    for {
        conn, err := lis.Accept()
        if err != nil {
            continue
        }

        go distribute(conn, backends)
    }
}

func distribute(conn net.Conn, backends []string) {
    defer conn.Close()

    backs := make([]net.Conn, 0, len(backends))
    for _, backend := range backends {
        back, err := net.DialTimeout("tcp", backend, 3*time.Second)
        if err != nil {
            continue
        }
        backs = append(backs, back)
    }
    if len(backs) == 0 {
        return
    }

    pipe := &pipeline{conn: conn, backs: backs}

    pipe.serve()
}

type pipeline struct {
    conn  net.Conn
    backs []net.Conn
}

func (p *pipeline) serve() {
    defer p.close()

    // 第一个连接为主连接，需要将数据响应给客户端
    master := p.backs[0]
    go p.readFrom(master)

    // 虽然不关心剩余的后端连接响应数据，但是应该将数据读出并丢弃
    discards := p.backs[1:]
    if len(discards) != 0 {
        go p.discard(discards)
    }

    buf := make([]byte, 1024)
    for {
        n, err := p.conn.Read(buf)
        if err != nil {
            break
        }

        // 读出客户端的数据，并分发给所有后端连接
        if hasErr := p.writeAll(buf[:n]); hasErr && len(p.backs) == 0 {
            break
        }
    }
}

func (p *pipeline) readFrom(r io.Reader) {
    _, _ = io.Copy(p.conn, r)
    _ = p.conn.Close()
}

func (*pipeline) discard(rds []net.Conn) {
    for _, rd := range rds {
        go io.Copy(io.Discard, rd)
    }
}

func (p *pipeline) writeAll(data []byte) bool {
    var hasErr bool
    for i, back := range p.backs {
        if _, err := back.Write(data); err == nil {
            continue
        }
        // 写入数据出现错误，说明后端连接异常，将这个连接剔除列表
        p.backs = append(p.backs[:i], p.backs[i+1:]...)
        hasErr = true
        _ = back.Close() // 保险起见，剔除后再主动执行 Close() 方法
    }

    return hasErr
}

func (p *pipeline) close() {
    _ = p.conn.Close()
    for _, back := range p.backs {
        _ = back.Close()
    }
}