From 8c6a74561f761fe65842aec73ae60ca64308c80d Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Mon, 7 Dec 2020 18:10:08 +0800
Subject: [PATCH] lru cache cert

---
 README.md     |  6 ++++--
 cert/cert.go  | 22 +++++++++++++++++++++-
 go.mod        |  6 +++++-
 go.sum        | 20 ++++++++++++++++++++
 proxy/mitm.go |  2 +-
 5 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 2d47003..db35648 100644
--- a/README.md
+++ b/README.md
@@ -6,7 +6,9 @@
 
 - [x] http handler
 - [x] http connect
-- [ ] cert
-- [ ] https handler
+- [x] cert
+- [x] https handler
 - [ ] http2
 - [ ] logger
+- [ ] 经内存转发 https 流量
+- [ ] 忽略某些错误例如：broken pipe, reset by peer, timeout
diff --git a/cert/cert.go b/cert/cert.go
index a6f06e1..3a10960 100644
--- a/cert/cert.go
+++ b/cert/cert.go
@@ -16,6 +16,8 @@ import (
 	"os"
 	"path/filepath"
 	"time"
+
+	"github.com/golang/groupcache/lru"
 )
 
 // reference
@@ -28,6 +30,7 @@ type CA struct {
 	rsa.PrivateKey
 	RootCert  x509.Certificate
 	StorePath string
+	cache     *lru.Cache
 }
 
 func NewCA(path string) (*CA, error) {
@@ -36,7 +39,10 @@ func NewCA(path string) (*CA, error) {
 		return nil, err
 	}
 
-	ca := &CA{StorePath: storePath}
+	ca := &CA{
+		StorePath: storePath,
+		cache:     lru.New(100),
+	}
 
 	if err := ca.load(); err != nil {
 		if err != caErrNotFound {
@@ -231,6 +237,20 @@ func (ca *CA) saveCert() error {
 	return ca.saveCertTo(file)
 }
 
+func (ca *CA) GetCert(commonName string) (*tls.Certificate, error) {
+	if val, ok := ca.cache.Get(commonName); ok {
+		return val.(*tls.Certificate), nil
+	}
+
+	cert, err := ca.DummyCert(commonName)
+	if err != nil {
+		return cert, err
+	}
+
+	ca.cache.Add(commonName, cert)
+	return cert, err
+}
+
 // TODO: 是否应该支持多个 SubjectAltName
 func (ca *CA) DummyCert(commonName string) (*tls.Certificate, error) {
 	template := &x509.Certificate{
diff --git a/go.mod b/go.mod
index 7ca4d73..296778d 100644
--- a/go.mod
+++ b/go.mod
@@ -2,4 +2,8 @@ module github.com/lqqyt2423/go-mitmproxy
 
 go 1.15
 
-require github.com/joho/godotenv v1.3.0
+require (
+	github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e
+	github.com/golang/protobuf v1.4.3 // indirect
+	github.com/joho/godotenv v1.3.0
+)
diff --git a/go.sum b/go.sum
index ead7071..647b21d 100644
--- a/go.sum
+++ b/go.sum
@@ -1,2 +1,22 @@
+github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e h1:1r7pUrabqp18hOBcwBwiTsbnFeTZHV9eER/QT5JVZxY=
+github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
+github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8=
+github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA=
+github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs=
+github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w=
+github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0=
+github.com/golang/protobuf v1.4.3 h1:JjCZWpVbqXDqFVmTfYWEVTMIYrL/NPdPSCHPJ0T/raM=
+github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI=
+github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
+github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
+github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/joho/godotenv v1.3.0 h1:Zjp+RcGpHhGlrMbJzXTrZZPrWj+1vfm90La1wgB6Bhc=
 github.com/joho/godotenv v1.3.0/go.mod h1:7hK45KPybAkOC6peb+G5yklZfMxEjkZhHbwpqxOKXbg=
+golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8=
+google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0=
+google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM=
+google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE=
+google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo=
+google.golang.org/protobuf v1.23.0 h1:4MY060fB1DLGMB/7MBTLnwQUY6+F09GEiz6SsrNqyzM=
+google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU=
diff --git a/proxy/mitm.go b/proxy/mitm.go
index 08d131f..616c4a5 100644
--- a/proxy/mitm.go
+++ b/proxy/mitm.go
@@ -53,7 +53,7 @@ func NewMitmServer(proxy *Proxy) (Mitm, error) {
 		TLSConfig: &tls.Config{
 			GetCertificate: func(chi *tls.ClientHelloInfo) (*tls.Certificate, error) {
 				// log.Printf("MitmServer GetCertificate ServerName: %v\n", chi.ServerName)
-				return ca.DummyCert(chi.ServerName)
+				return ca.GetCert(chi.ServerName)
 			},
 		},
 	}