Go 学习笔记（13）：开发一个简单的端口转发程序 ​

package main

import (
	"context"
	"fmt"
	"gopkg.in/yaml.v2"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
)

// Config 结构定义了配置文件的结构
type Config struct {
	PortForwards []PortForward `yaml:"port_forwards"`
}

// PortForward 结构定义了单个端口转发的配置
type PortForward struct {
	LocalPort    int    `yaml:"local_port"`
	RemoteAddr   string `yaml:"remote_addr"`
	RemotePort   int    `yaml:"remote_port"`
	ProtocolType string `yaml:"protocol_type"`
}

func main() {
	// 读取配置文件
	config, err := readConfig("config.yaml")
	if err != nil {
		log.Fatalf("Failed to read config file: %v", err)
	}

	// 创建上下文
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// 启动端口转发服务
	var wg sync.WaitGroup
	for _, pf := range config.PortForwards {
		wg.Add(1)
		go func(pf PortForward) {
			defer wg.Done()
			startPortForward(ctx, cancel, pf)
		}(pf)
	}
	wg.Wait()
}

// readConfig 从配置文件中读取配置信息
func readConfig(filename string) (*Config, error) {
	data, err := os.ReadFile(filename)
	if err != nil {
		return nil, fmt.Errorf("failed to read file: %v", err)
	}

	var config Config
	err = yaml.Unmarshal(data, &config)
	if err != nil {
		return nil, fmt.Errorf("failed to unmarshal config: %v", err)
	}

	return &config, nil
}

// startPortForward 启动端口转发服务
func startPortForward(ctx context.Context, cancel context.CancelFunc, pf PortForward) {
	sourcePort := fmt.Sprintf(":%v", pf.LocalPort)
	destinationAddress := fmt.Sprintf("%s:%v", pf.RemoteAddr, pf.RemotePort)

	listener, err := net.Listen("tcp", sourcePort)
	if err != nil {
		log.Fatalf("Failed to listen on %s: %v", sourcePort, err)
	}
	defer func(listener net.Listener) {
		_ = listener.Close()
	}(listener)
	log.Printf("PID %d: Listening on %s and forwarding to %s", os.Getpid(), sourcePort, destinationAddress)

	// 监听系统信号
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		sig := <-sigChan
		log.Printf("PID %d: Received signal: %s. Shutting down.", os.Getpid(), sig)
		cancel() // 取消上下文，停止新的连接接受
		_ = listener.Close()
	}()

	for {
		select {
		case <-ctx.Done():
			log.Printf("PID %d: Context cancelled, shutting down listener on %s", os.Getpid(), sourcePort)
			return
		default:
			clientConn, err := listener.Accept()
			if err != nil {
				if strings.Contains(err.Error(), net.ErrClosed.Error()) {
					log.Printf("PID %d: Temporary accept error: %v", os.Getpid(), err)
					return // 退出循环，避免持续打印错误日志
				}
				log.Printf("PID %d: Failed to accept connection: %v", os.Getpid(), err)
				break
			}
			log.Printf("PID %d: Accepted connection from %s", os.Getpid(), clientConn.RemoteAddr())
			go handleConnection(ctx, clientConn, destinationAddress)
		}
	}
}

// handleConnection 处理连接
func handleConnection(ctx context.Context, clientConn net.Conn, destinationAddress string) {
	defer func(clientConn net.Conn) {
		_ = clientConn.Close()
	}(clientConn)
	log.Printf("PID %d: Handling connection to %s", os.Getpid(), destinationAddress)

	serverConn, err := net.Dial("tcp", destinationAddress)
	if err != nil {
		log.Printf("PID %d: Failed to connect to destination %s: %v", os.Getpid(), destinationAddress, err)
		return
	}
	defer func(serverConn net.Conn) {
		_ = serverConn.Close()
	}(serverConn)

	doneChan := make(chan struct{})

	go copyData(ctx, clientConn, serverConn, doneChan)
	go copyData(ctx, serverConn, clientConn, doneChan)

	select {
	case <-ctx.Done():
		log.Printf("PID %d: Context cancelled, closing connections", os.Getpid())
	case <-doneChan:
		log.Printf("PID %d: Data transfer completed, closing connections", os.Getpid())
	}
}

// copyData 复制数据
func copyData(ctx context.Context, dst net.Conn, src net.Conn, doneChan chan struct{}) {
	log.Printf("PID %d: Starting data copy from %s to %s", os.Getpid(), src.RemoteAddr(), dst.RemoteAddr())
	_, err := io.Copy(dst, src)
	if err != nil {
		select {
		case <-ctx.Done():
			log.Printf("PID %d: Context cancelled, stopping data copy from %s to %s", os.Getpid(), src.RemoteAddr(), dst.RemoteAddr())
			return
		default:
			if strings.Contains(err.Error(), "use of closed network connection") {
				log.Printf("PID %d: Connection closed during data copy from %s to %s", os.Getpid(), src.RemoteAddr(), dst.RemoteAddr())
			} else {
				log.Printf("PID %d: Error copying data from %s to %s: %v", os.Getpid(), src.RemoteAddr(), dst.RemoteAddr(), err)
			}
		}
	}
	log.Printf("PID %d: Completed data copy from %s to %s", os.Getpid(), src.RemoteAddr(), dst.RemoteAddr())
	doneChan <- struct{}{}
}

package main

import (
    "context"
    "fmt"
    "time"
)

func main() {
    ctx := context.Background()
    fmt.Println(ctx)
    
    todoCtx := context.TODO()
    fmt.Println(todoCtx)
}

package main

import (
    "context"
    "fmt"
    "time"
)

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()

    go func() {
        // 模拟工作
        time.Sleep(2 * time.Second)
        cancel()  // 取消上下文
    }()

    select {
    case <-ctx.Done():
        fmt.Println("工作取消:", ctx.Err())
    }
}

package main

import (
    "context"
    "fmt"
    "time"
)

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    defer cancel()

    select {
    case <-time.After(5 * time.Second):
        fmt.Println("操作完成")
    case <-ctx.Done():
        fmt.Println("操作超时:", ctx.Err())
    }
}

package main

import (
    "context"
    "fmt"
    "time"
)

func main() {
    deadline := time.Now().Add(3 * time.Second)
    ctx, cancel := context.WithDeadline(context.Background(), deadline)
    defer cancel()

    select {
    case <-time.After(5 * time.Second):
        fmt.Println("操作完成")
    case <-ctx.Done():
        fmt.Println("达到截止日期:", ctx.Err())
    }
}

package main

import (
    "context"
    "fmt"
    "time"
)

func worker(ctx context.Context, id int) {
    for {
        select {
        case <-ctx.Done():
            fmt.Printf("工人 %d: 接收到取消信号\n", id)
            return
        default:
            fmt.Printf("工人 %d: 正在工作\n", id)
            time.Sleep(500 * time.Millisecond)
        }
    }
}

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()

    for i := 1; i <= 3; i++ {
        go worker(ctx, i)
    }

    // 等待所有工作完成或超时
    <-ctx.Done()
    fmt.Println("主程序: 取消所有工人")
}

package main

import (
    "context"
    "fmt"
    "net/http"
    "time"
)

func handler(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()
    fmt.Println("处理请求")
    select {
    case <-time.After(5 * time.Second):
        fmt.Fprintln(w, "请求处理完成")
    case <-ctx.Done():
        fmt.Println("请求取消:", ctx.Err())
        http.Error(w, "请求取消", http.StatusRequestTimeout)
    }
}

func main() {
    http.HandleFunc("/", handler)
    http.ListenAndServe(":8080", nil)
}

package main

import (
    "fmt"
    "os"
    "os/signal"
    "syscall"
)

func main() {
    // 创建一个通道用于接收信号
    sigs := make(chan os.Signal, 1)
    // 注册信号
    signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

    // 使用一个goroutine等待信号
    go func() {
        sig := <-sigs
        fmt.Println("收到信号:", sig)
        // 在这里可以处理信号，例如进行清理操作
        os.Exit(0)
    }()

    // 模拟程序运行
    fmt.Println("程序运行中，按Ctrl+C终止")
    select {}
}

sigs := make(chan os.Signal, 1)

signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)

signal.Stop(sigs)