在Go中构建应用级IP防火墙机制

2023 年 11 月 26 日

141 次浏览

3248字数

构建应用级IP防火墙机制可以通过在应用层面进行请求过滤来实现。在Go语言中，这通常意味着编写中间件，该中间件会检查进入应用程序的每个请求的IP地址，并决定允许或拒绝该请求。以下是一个实现应用级IP防火墙的示例。

package main

import (
    "fmt"
    "log"
    "net"
    "net/http"
)

// IPFilterMiddleware 创建一个新的中间件实例来管理IP访问策略。
type IPFilterMiddleware struct {
    // 黑名单模式或白名单模式，默认为黑名单。
    blacklist bool
    // 允许或阻止的IP列表
    ips map[string]bool
}

// NewIPFilterMiddleware 初始化IP过滤中间件。
func NewIPFilterMiddleware(blacklistMode bool) *IPFilterMiddleware {
    return &IPFilterMiddleware{
        blacklist: blacklistMode,
        ips:       make(map[string]bool),
    }
}

// AddIP 添加IP到列表中。
func (m *IPFilterMiddleware) AddIP(ip string) {
    m.ips[ip] = true
}

// RemoveIP 从列表中移除IP。
func (m *IPFilterMiddleware) RemoveIP(ip string) {
    delete(m.ips, ip)
}

// Middleware 作为HTTP中间件函数。
func (m *IPFilterMiddleware) Middleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        ip, _, err := net.SplitHostPort(r.RemoteAddr)
        if err != nil {
            http.Error(w, "Internal Server Error", http.StatusInternalServerError)
            return
        }

        if m.blacklist {
            // 黑名单模式：如果IP在黑名单中，则拒绝。
            if m.ips[ip] {
                http.Error(w, "Forbidden", http.StatusForbidden)
                return
            }
        } else {
            // 白名单模式：如果IP不在白名单中，则拒绝。
            if !m.ips[ip] {
                http.Error(w, "Forbidden", http.StatusForbidden)
                return
            }
        }

        // 继续执行其他中间件或最终处理程序
        next.ServeHTTP(w, r)
    })
}

func main() {
    // 初始化IP过滤中间件，true表示黑名单模式
    ipFilter := NewIPFilterMiddleware(true)

    // 添加一些IP到黑名单
    ipFilter.AddIP("192.168.1.1")
    ipFilter.AddIP("10.0.0.1")

    // 设置HTTP服务器
    mux := http.NewServeMux()
    mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
        fmt.Fprintln(w, "Hello World!")
    })

    // 应用IP过滤中间件
    filteredHandler := ipFilter.Middleware(mux)

    log.Println("Server is running on :8080")
    if err := http.ListenAndServe(":8080", filteredHandler); err != nil {
        log.Fatalf("Unable to start server: %s\n", err)
    }
}

在这个示例中，我们创建了一个 IPFilterMiddleware结构体，这个结构体能够让你定义一个IP列表以及是将这个列表作为黑名单还是白名单进行处理。IP列表是作为一个 map来存储IP，其值表示IP的存在与否。

在 Middleware函数中，我们定义了一个闭包，这个闭包会被每个进入的HTTP请求所调用。闭包内部我们检查请求的IP地址，根据是黑名单还是白名单模式来决定采取行动。如果是黑名单模式，并且请求IP在黑名单中，或者是白名单模式，并且请求IP不在白名单中，这个请求将被拒绝，并返回HTTP状态码403（Forbidden）。否则，请求会被传递到链中的下一个处理程序（可能是另一个中间件或者是最终的处理函数）。

在 main函数中，我们设置了服务器和中间件，并启动HTTP服务器。

使用Go构建应用级别的IP防火墙机制不仅能够为你的应用程序增加一层额外的安全性，还能够通过自定义中间件的方式让你有更多控制力，来决定哪些客户端可以或不可以访问你的服务。