TCP 与 UDP 编程

38.1 网络编程基础

TCP vs UDP

Go net 包

Go 的 net 包提供了网络 I/O 的接口，包括 TCP、UDP、Unix 域套接字等。

38.2 TCP 服务器

基本 TCP 服务器

package main

import (
    "bufio"
    "fmt"
    "net"
)

func main() {
    // 监听端口
    listener, err := net.Listen("tcp", ":8080")
    if err != nil {
        fmt.Println("监听失败:", err)
        return
    }
    defer listener.Close()

    fmt.Println("TCP 服务器启动在 :8080")

    for {
        // 接受连接
        conn, err := listener.Accept()
        if err != nil {
            fmt.Println("接受连接失败:", err)
            continue
        }

        // 处理连接
        go handleConnection(conn)
    }
}

func handleConnection(conn net.Conn) {
    defer conn.Close()

    remoteAddr := conn.RemoteAddr().String()
    fmt.Printf("客户端连接: %s\n", remoteAddr)

    scanner := bufio.NewScanner(conn)
    for scanner.Scan() {
        text := scanner.Text()
        fmt.Printf("收到: %s\n", text)

        // 回显
        conn.Write([]byte("Echo: " + text + "\n"))
    }

    fmt.Printf("客户端断开: %s\n", remoteAddr)
}

并发 TCP 服务器

package main

import (
    "bufio"
    "fmt"
    "net"
    "sync"
)

type Server struct {
    clients map[net.Conn]bool
    mu      sync.RWMutex
}

func NewServer() *Server {
    return &Server{
        clients: make(map[net.Conn]bool),
    }
}

func (s *Server) addClient(conn net.Conn) {
    s.mu.Lock()
    s.clients[conn] = true
    s.mu.Unlock()
}

func (s *Server) removeClient(conn net.Conn) {
    s.mu.Lock()
    delete(s.clients, conn)
    s.mu.Unlock()
}

func (s *Server) broadcast(message string, sender net.Conn) {
    s.mu.RLock()
    defer s.mu.RUnlock()

    for conn := range s.clients {
        if conn != sender {
            conn.Write([]byte(message + "\n"))
        }
    }
}

func (s *Server) handleConnection(conn net.Conn) {
    defer conn.Close()
    defer s.removeClient(conn)

    s.addClient(conn)
    fmt.Printf("客户端连接: %s\n", conn.RemoteAddr())

    scanner := bufio.NewScanner(conn)
    for scanner.Scan() {
        text := scanner.Text()
        message := fmt.Sprintf("%s: %s", conn.RemoteAddr(), text)
        fmt.Println(message)
        s.broadcast(message, conn)
    }
}

func main() {
    server := NewServer()

    listener, _ := net.Listen("tcp", ":8080")
    defer listener.Close()

    fmt.Println("聊天服务器启动在 :8080")

    for {
        conn, _ := listener.Accept()
        go server.handleConnection(conn)
    }
}

38.3 TCP 客户端

基本 TCP 客户端

package main

import (
    "bufio"
    "fmt"
    "net"
    "os"
)

func main() {
    // 连接服务器
    conn, err := net.Dial("tcp", "localhost:8080")
    if err != nil {
        fmt.Println("连接失败:", err)
        return
    }
    defer conn.Close()

    fmt.Println("已连接到服务器")

    // 读取服务器响应
    go func() {
        scanner := bufio.NewScanner(conn)
        for scanner.Scan() {
            fmt.Println("服务器:", scanner.Text())
        }
    }()

    // 发送消息
    scanner := bufio.NewScanner(os.Stdin)
    for scanner.Scan() {
        text := scanner.Text()
        conn.Write([]byte(text + "\n"))
    }
}

带超时的客户端

package main

import (
    "bufio"
    "fmt"
    "net"
    "time"
)

func main() {
    // 带超时的连接
    conn, err := net.DialTimeout("tcp", "localhost:8080", 5*time.Second)
    if err != nil {
        fmt.Println("连接失败:", err)
        return
    }
    defer conn.Close()

    // 设置读写超时
    conn.SetDeadline(time.Now().Add(10 * time.Second))

    // 发送数据
    conn.Write([]byte("Hello, Server!\n"))

    // 读取响应
    scanner := bufio.NewScanner(conn)
    if scanner.Scan() {
        fmt.Println("服务器:", scanner.Text())
    }
}

38.4 UDP 服务器

基本 UDP 服务器

package main

import (
    "fmt"
    "net"
)

func main() {
    // 监听 UDP 端口
    addr, _ := net.ResolveUDPAddr("udp", ":8080")
    conn, err := net.ListenUDP("udp", addr)
    if err != nil {
        fmt.Println("监听失败:", err)
        return
    }
    defer conn.Close()

    fmt.Println("UDP 服务器启动在 :8080")

    buf := make([]byte, 1024)

    for {
        // 接收数据
        n, remoteAddr, err := conn.ReadFromUDP(buf)
        if err != nil {
            fmt.Println("读取失败:", err)
            continue
        }

        message := string(buf[:n])
        fmt.Printf("收到来自 %s: %s\n", remoteAddr, message)

        // 发送响应
        response := "Echo: " + message
        conn.WriteToUDP([]byte(response), remoteAddr)
    }
}

并发 UDP 服务器

package main

import (
    "fmt"
    "net"
    "sync"
)

func main() {
    addr, _ := net.ResolveUDPAddr("udp", ":8080")
    conn, _ := net.ListenUDP("udp", addr)
    defer conn.Close()

    fmt.Println("UDP 服务器启动在 :8080")

    var wg sync.WaitGroup
    buf := make([]byte, 1024)

    for {
        n, remoteAddr, err := conn.ReadFromUDP(buf)
        if err != nil {
            continue
        }

        wg.Add(1)
        go func(data []byte, addr *net.UDPAddr) {
            defer wg.Done()
            handleUDPMessage(conn, data, addr)
        }(buf[:n], remoteAddr)
    }
}

func handleUDPMessage(conn *net.UDPConn, data []byte, addr *net.UDPAddr) {
    message := string(data)
    fmt.Printf("处理来自 %s: %s\n", addr, message)

    response := "Processed: " + message
    conn.WriteToUDP([]byte(response), addr)
}

38.5 UDP 客户端

基本 UDP 客户端

package main

import (
    "bufio"
    "fmt"
    "net"
    "os"
)

func main() {
    // 解析服务器地址
    serverAddr, err := net.ResolveUDPAddr("udp", "localhost:8080")
    if err != nil {
        fmt.Println("解析地址失败:", err)
        return
    }

    // 创建本地 UDP 套接字
    localAddr, _ := net.ResolveUDPAddr("udp", ":0")
    conn, err := net.DialUDP("udp", localAddr, serverAddr)
    if err != nil {
        fmt.Println("连接失败:", err)
        return
    }
    defer conn.Close()

    fmt.Println("UDP 客户端已启动")

    // 读取用户输入并发送
    scanner := bufio.NewScanner(os.Stdin)
    for scanner.Scan() {
        text := scanner.Text()

        // 发送数据
        conn.Write([]byte(text))

        // 接收响应
        buf := make([]byte, 1024)
        n, _ := conn.Read(buf)
        fmt.Println("服务器:", string(buf[:n]))
    }
}

简单 UDP 客户端

package main

import (
    "fmt"
    "net"
)

func main() {
    conn, _ := net.Dial("udp", "localhost:8080")
    defer conn.Close()

    // 发送数据
    conn.Write([]byte("Hello, UDP Server!"))

    // 接收响应
    buf := make([]byte, 1024)
    n, _ := conn.Read(buf)
    fmt.Println("响应:", string(buf[:n]))
}

38.6 TCP 文件传输

文件服务器

package main

import (
    "fmt"
    "io"
    "net"
    "os"
)

func main() {
    listener, _ := net.Listen("tcp", ":8080")
    defer listener.Close()

    fmt.Println("文件服务器启动在 :8080")

    for {
        conn, _ := listener.Accept()
        go handleFileTransfer(conn)
    }
}

func handleFileTransfer(conn net.Conn) {
    defer conn.Close()

    // 接收文件名
    buf := make([]byte, 1024)
    n, _ := conn.Read(buf)
    filename := string(buf[:n])

    // 发送确认
    conn.Write([]byte("OK"))

    // 创建文件
    file, err := os.Create("received_" + filename)
    if err != nil {
        fmt.Println("创建文件失败:", err)
        return
    }
    defer file.Close()

    // 接收文件内容
    _, err = io.Copy(file, conn)
    if err != nil {
        fmt.Println("接收文件失败:", err)
        return
    }

    fmt.Printf("文件 %s 接收完成\n", filename)
}

文件客户端

package main

import (
    "fmt"
    "io"
    "net"
    "os"
)

func main() {
    if len(os.Args) < 2 {
        fmt.Println("用法: client <文件名>")
        return
    }
    filename := os.Args[1]

    // 连接服务器
    conn, _ := net.Dial("tcp", "localhost:8080")
    defer conn.Close()

    // 发送文件名
    conn.Write([]byte(filename))

    // 等待确认
    buf := make([]byte, 2)
    conn.Read(buf)

    // 发送文件内容
    file, _ := os.Open(filename)
    defer file.Close()

    io.Copy(conn, file)

    fmt.Println("文件发送完成")
}

38.7 自定义协议

定义协议格式

package main

import (
    "encoding/binary"
    "fmt"
    "net"
)

// 协议格式: | 长度(4字节) | 类型(1字节) | 数据 |

const (
    TypeMessage = 1
    TypeFile    = 2
    TypeCommand = 3
)

type Packet struct {
    Type uint8
    Data []byte
}

func writePacket(conn net.Conn, p *Packet) error {
    // 写入长度
    length := uint32(len(p.Data) + 1)
    binary.Write(conn, binary.BigEndian, length)

    // 写入类型
    binary.Write(conn, binary.BigEndian, p.Type)

    // 写入数据
    _, err := conn.Write(p.Data)
    return err
}

func readPacket(conn net.Conn) (*Packet, error) {
    // 读取长度
    var length uint32
    err := binary.Read(conn, binary.BigEndian, &length)
    if err != nil {
        return nil, err
    }

    // 读取类型
    var pType uint8
    binary.Read(conn, binary.BigEndian, &pType)

    // 读取数据
    data := make([]byte, length-1)
    _, err = conn.Read(data)
    if err != nil {
        return nil, err
    }

    return &Packet{Type: pType, Data: data}, nil
}

func main() {
    listener, _ := net.Listen("tcp", ":8080")
    defer listener.Close()

    fmt.Println("协议服务器启动在 :8080")

    for {
        conn, _ := listener.Accept()
        go func() {
            defer conn.Close()

            for {
                packet, err := readPacket(conn)
                if err != nil {
                    return
                }

                switch packet.Type {
                case TypeMessage:
                    fmt.Printf("消息: %s\n", packet.Data)
                case TypeCommand:
                    fmt.Printf("命令: %s\n", packet.Data)
                }

                // 发送响应
                writePacket(conn, &Packet{
                    Type: TypeMessage,
                    Data: []byte("收到"),
                })
            }
        }()
    }
}

38.8 心跳机制

带心跳的连接

package main

import (
    "fmt"
    "net"
    "time"
)

type HeartbeatConn struct {
    conn        net.Conn
    timeout     time.Duration
    lastActive  time.Time
}

func NewHeartbeatConn(conn net.Conn, timeout time.Duration) *HeartbeatConn {
    return &HeartbeatConn{
        conn:       conn,
        timeout:    timeout,
        lastActive: time.Now(),
    }
}

func (hc *HeartbeatConn) Read(b []byte) (n int, err error) {
    n, err = hc.conn.Read(b)
    hc.lastActive = time.Now()
    return
}

func (hc *HeartbeatConn) Write(b []byte) (n int, err error) {
    n, err = hc.conn.Write(b)
    hc.lastActive = time.Now()
    return
}

func (hc *HeartbeatConn) Close() error {
    return hc.conn.Close()
}

func (hc *HeartbeatConn) CheckTimeout() bool {
    return time.Since(hc.lastActive) > hc.timeout
}

func main() {
    listener, _ := net.Listen("tcp", ":8080")
    defer listener.Close()

    fmt.Println("心跳服务器启动在 :8080")

    for {
        conn, _ := listener.Accept()
        hc := NewHeartbeatConn(conn, 30*time.Second)

        go func() {
            defer hc.Close()

            ticker := time.NewTicker(10 * time.Second)
            defer ticker.Stop()

            buf := make([]byte, 1024)

            for {
                select {
                case <-ticker.C:
                    if hc.CheckTimeout() {
                        fmt.Println("连接超时，关闭")
                        return
                    }
                    // 发送心跳
                    hc.Write([]byte("PING\n"))
                default:
                    hc.conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
                    n, err := hc.Read(buf)
                    if err == nil {
                        message := string(buf[:n])
                        if message == "PONG\n" {
                            continue
                        }
                        fmt.Printf("收到: %s", message)
                        hc.Write([]byte("Echo: " + message))
                    }
                }
            }
        }()
    }
}

38.9 小结

本章详细介绍了 Go 语言的 TCP 和 UDP 编程：

1. TCP 服务器：使用 net.Listen 创建 TCP 服务器
2. TCP 客户端：使用 net.Dial 连接 TCP 服务器
3. UDP 服务器：使用 net.ListenUDP 创建 UDP 服务器
4. UDP 客户端：使用 net.DialUDP 连接 UDP 服务器
5. 并发处理：为每个连接创建 Goroutine
6. 文件传输：基于 TCP 实现文件传输
7. 自定义协议：定义自己的通信协议
8. 心跳机制：保持连接活跃

网络编程是构建分布式系统的基础，掌握 TCP 和 UDP 编程能让你构建各种网络应用。在下一章中，我们将学习测试与调试。