go-mitmproxy/proxy/helper.go

package proxy

import (
	"bytes"
	"io"
	"os"
	"strings"
	"sync"

	_log "github.com/sirupsen/logrus"
)

var NormalErrMsgs []string = []string{
	"read: connection reset by peer",
	"write: broken pipe",
	"i/o timeout",
	"net/http: TLS handshake timeout",
	"io: read/write on closed pipe",
	"connect: connection refused",
	"connect: connection reset by peer",
}

// 仅打印预料之外的错误信息
func LogErr(log *_log.Entry, err error) (loged bool) {
	msg := err.Error()

	for _, str := range NormalErrMsgs {
		if strings.Contains(msg, str) {
			log.Debug(err)
			return
		}
	}

	log.Error(err)
	loged = true
	return
}

// 转发流量
// Read a => Write b
// Read b => Write a
func Transfer(log *_log.Entry, a, b io.ReadWriteCloser) {
	done := make(chan struct{})
	defer close(done)

	forward := func(dst io.WriteCloser, src io.Reader, ec chan<- error) {
		_, err := io.Copy(dst, src)
		if err != nil {
			select {
			case <-done:
				return
			case ec <- err:
				return
			}
		}

		err = dst.Close()
		select {
		case <-done:
			return
		case ec <- err:
			return
		}
	}

	errChan := make(chan error)
	go forward(a, b, errChan)
	go forward(b, a, errChan)

	for i := 0; i < 2; i++ {
		if err := <-errChan; err != nil {
			LogErr(log, err)
			return // 如果有错误，直接返回
		}
	}
}

// 尝试将 Reader 读取至 buffer 中
// 如果未达到 limit，则成功读取进入 buffer
// 否则 buffer 返回 nil，且返回新 Reader，状态为未读取前
func ReaderToBuffer(r io.Reader, limit int64) ([]byte, io.Reader, error) {
	buf := bytes.NewBuffer(make([]byte, 0))
	lr := io.LimitReader(r, limit)

	_, err := io.Copy(buf, lr)
	if err != nil {
		return nil, nil, err
	}

	// 达到上限
	if int64(buf.Len()) == limit {
		// 返回新的 Reader
		return nil, io.MultiReader(bytes.NewBuffer(buf.Bytes()), r), nil
	}

	// 返回 buffer
	return buf.Bytes(), nil, nil
}

// Wireshark 解析 https 设置
var tlsKeyLogWriter io.Writer
var tlsKeyLogOnce sync.Once

func GetTlsKeyLogWriter() io.Writer {
	tlsKeyLogOnce.Do(func() {
		logfile := os.Getenv("SSLKEYLOGFILE")
		if logfile == "" {
			return
		}

		writer, err := os.OpenFile(logfile, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
		if err != nil {
			log.WithField("in", "GetTlsKeyLogWriter").Debug(err)
			return
		}

		tlsKeyLogWriter = writer
	})
	return tlsKeyLogWriter
}