基于 Gin 的 HTTP 代理上网行为记录 demo

前言: 前端时间写了好几篇使用 Gin 框架来做 HTTP 代理 demo 的文章，然后就想着做一个记录上网行为的小工具，就是简单记录看看平时访问了什么网站（基于隧道代理的，不是中间人代理，所以只能记录去了哪里，不能记录干了什么）。不过因为编译问题一直没有解决，我又不想重新在 Windows 上安装 Golang 的开发环境，所以就把它搁置了。最近正好把那个交叉编译的问题解决了，所以就把这个博客也发出来吧。

一、代码

主要的代码还是前几篇博客已经介绍过的了，我只是加了一个写入 sqlite3 的功能。因为也没有必要有一条一条的写入，所以就是每 100 条写入一次了。程序运行会在当前目录生成一个 net_access.db 的文件，这个就是 sqlite 的数据库文件。这里只有一个记录访问了什么网站的表，本来还想着做一个 block_list 的表，阻止访问某些网站的（就是不给它建立连接），不过放弃了。

package main

import (
	"bufio"
	"io"
	"log"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

type ReqRecord struct {
	ID        uint      `gorm:"primaryKey,column:id"`
	CreatedAt time.Time `gorm:"column:time"`
	Schema    string    `gorm:"column:schema"` // http/https
	Host      string    `gorm:"column:host"`
	Port      string    `gorm:"column:port"`
}

const (
	RPOXY_SERVER  = "CrazyDragonHttpProxy"                                                             // it is just a kidding, but Only HTTP!
	TUNNEL_PACKET = "HTTP/1.1 200 Connection Established\r\nProxy-agent: CrazyDragonHttpProxy\r\n\r\n" // Don'e USE `` to surround a protocl strng, DAMN!!!
	BUFFER_SIZE   = 100
)

var (
	proxyHttpClient = http.DefaultClient
	DBConn          *gorm.DB
	recordChan      chan *ReqRecord
)

func init() {
	var err error
	DBConn, err = gorm.Open(sqlite.Open("net_access.db"), &gorm.Config{})
	if err != nil {
		panic("failed to connect database")
	}
	err = DBConn.AutoMigrate(&ReqRecord{})
	if err != nil {
		panic(err)
	}

	recordChan = make(chan *ReqRecord, BUFFER_SIZE) // 缓冲区大小为 100
}

func main() {
	go write2DB()
	r := gin.Default()
	r.NoRoute(routeProxy)   // NO Route is every Route!!!
	r.Run("127.0.0.1:8888")
}

// Then I can process all routes
func routeProxy(c *gin.Context) {
	req := c.Request
	go recordReq(req) // 记录日志
	if req.Method == http.MethodConnect {
		httpsProxy(c, req) // create http tunnel to process https
	} else {
		httpProxy(c, req) // process plain http
	}
}

func httpsProxy(c *gin.Context, req *http.Request) {
	// established connect tunnel
	address := req.URL.Host // it contains the port
	tunnelConn, err := net.DialTimeout("tcp", address, 10*time.Second)
	if err != nil {
		log.Println(err)
		return
	}
	log.Printf("try to established Connect Tunnel to: %s has been successfully.\n", address)
	tunnelrw := bufio.NewReadWriter(bufio.NewReader(tunnelConn), bufio.NewWriter(tunnelConn))
	c.Status(200) // 不加也没有问题，只是说 https 响应都会变成 404.
	// And We need to take over the http connection, Then make it become a TCP connection.
	hj, ok := c.Writer.(http.Hijacker)
	if !ok {
		http.Error(c.Writer, "webserver doesn't support hijacking", http.StatusInternalServerError)
		return
	}
	clientConn, bufrw, err := hj.Hijack()
	if err != nil {
		http.Error(c.Writer, err.Error(), http.StatusInternalServerError)
		return
	}

	// 建立隧道之后，发送一个连接建立的响应（其实，只要状态码是 200 就可以了）
	if _, err = clientConn.Write([]byte(TUNNEL_PACKET)); err != nil {
		log.Printf("Response Failed: %v", err.Error())
	} else {
		log.Println("Response Success.")
	}

	// data flow direction: client <---> tunnel <---> server
	defer clientConn.Close()
	defer tunnelConn.Close()
	done := make(chan struct{})
	go transfer(bufrw, tunnelrw, done) // client --> proxy --> server
	go transfer(tunnelrw, bufrw, done) //server --> proxy --> client
	<-done
}

func httpProxy(c *gin.Context, req *http.Request) {
	req.RequestURI = "" // Must create a new Req or empty this.
	resp, err := proxyHttpClient.Do(req)
	if err != nil {
		log.Println(err)
		return
	}
	defer resp.Body.Close()
	c.Status(resp.StatusCode) // change the status code, default is 404 !!!
	for k, v := range resp.Header {
		c.Header(k, strings.Join(v, ",")) // write Header
	}
	c.Header("Server", RPOXY_SERVER) // haha, it just a kidding!!!
	io.Copy(c.Writer, resp.Body)     // and response data to client
}

// tunnel transfer data.
func transfer(from io.Reader, to io.Writer, ch chan<- struct{}) {
	io.Copy(to, from)
	ch <- struct{}{}
}

// Record http request
func recordReq(req *http.Request) {
	var schema = "HTTPS"
	if req.Method != http.MethodConnect {
		schema = "HTTP"
	}
	record := &ReqRecord{
		Schema: schema,
		Host:   req.Host,
		Port:   req.URL.Port(),
	}
	recordChan <- record
	log.Printf("Method: %s, Host: %s, URL: %s, Version: %s\n", req.Method, req.Host, req.URL.Path, req.Proto)
}

func write2DB() {
	log.Printf("Start to record request info...")
	records := make([]*ReqRecord, BUFFER_SIZE)
	idx := 0
	for {
		r := <-recordChan
		records[idx] = r
		idx++
		if idx == BUFFER_SIZE {
			// 批量写入
			DBConn.Create(&records)
			idx = 0 // reset position
			log.Printf("Write to DB successfully.")
		}
	}
}