email/dkim: add key pool

Implementor of this library can use the KeyPool.Get method to retrieve key
instead of LookupKey to minimize network traffic and process to decode and
parse public key.

This changes affect how we call LookupKey.

6 files changed, 266 insertions, 49 deletions

diff --git a/lib/email/dkim/dkim.go b/lib/email/dkim/dkim.go
index b1a63b59..46d13fa9 100644
--- a/lib/email/dkim/dkim.go
+++ b/lib/email/dkim/dkim.go
@@ -6,12 +6,24 @@ package dkim
 
 //
 // DefaultNameServers contains list of nameserver's IP addresses.
-// If its not empty, the public key lookup using DNS/TXT will use this values.
 //
-var DefaultNameServers []string // nolint: gochecknoglobals
+// If its not empty, the public key lookup using DNS/TXT will use this values;
+// otherwise it will try to use the system name servers.
+//
+var DefaultNameServers []string // nolint:gochecknoglobals
+
+//
+// DefaultKeyPool contains cached DKIM key.
+//
+// Implementor of this library can use the KeyPool.Get method to retrieve key
+// instead of LookupKey to minimize network traffic and process to decode and
+// parse public key.
+//
+var DefaultKeyPool = &KeyPool{ // nolint:gochecknoglobals
+	pool: make(map[string]*Key),
+}
 
-var ( // nolint: gochecknoglobals
-	sepColon      = []byte{':'}
-	sepVBar       = []byte{'|'}
-	dkimSubdomain = []byte("_domainkey")
+var ( // nolint:gochecknoglobals
+	sepColon = []byte{':'}
+	sepVBar  = []byte{'|'}
 )
diff --git a/lib/email/dkim/key.go b/lib/email/dkim/key.go
index ea535560..422abaaa 100644
--- a/lib/email/dkim/key.go
+++ b/lib/email/dkim/key.go
@@ -10,6 +10,7 @@ import (
 	"crypto/x509"
 	"encoding/base64"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/shuLhan/share/lib/dns"
@@ -66,16 +67,15 @@ type Key struct {
 }
 
 //
-// LookupKey DKIM (public) key using specific query method, domain name
-// (SDID), and selector.
-// The sdid MUST NOT be empty, but the selector MAY be empty.
+// LookupKey DKIM (public) key using specific query method and DKIM domain
+// name (selector plus SDID).
 //
-func LookupKey(qmethod QueryMethod, sdid, selector []byte) (key *Key, err error) {
-	if len(sdid) == 0 {
-		return nil, fmt.Errorf("dkim: LookupKey: empty SDID")
+func LookupKey(qmethod QueryMethod, dname string) (key *Key, err error) {
+	if len(dname) == 0 {
+		return nil, nil
 	}
 	if qmethod.Type == QueryTypeDNS {
-		key, err = lookupDNS(qmethod.Option, sdid, selector)
+		key, err = lookupDNS(qmethod.Option, dname)
 	}
 	return key, err
 }
@@ -106,14 +106,14 @@ func ParseTXT(txt []byte, ttl uint32) (key *Key, err error) {
 	return key, nil
 }
 
-func lookupDNS(opt QueryOption, sdid, selector []byte) (key *Key, err error) {
+func lookupDNS(opt QueryOption, dname string) (key *Key, err error) {
 	if opt == QueryOptionTXT {
-		key, err = lookupDNSTXT(sdid, selector)
+		key, err = lookupDNSTXT(dname)
 	}
 	return key, err
 }
 
-func lookupDNSTXT(sdid, selector []byte) (key *Key, err error) {
+func lookupDNSTXT(dname string) (key *Key, err error) {
 	if dnsClientPool == nil {
 		err = newDNSClientPool()
 		if err != nil {
@@ -123,34 +123,25 @@ func lookupDNSTXT(sdid, selector []byte) (key *Key, err error) {
 
 	dnsClient := dnsClientPool.Get()
 
-	var bb bytes.Buffer
-	if len(selector) > 0 {
-		bb.Write(selector)
-		bb.WriteByte('.')
-	}
-	bb.Write(dkimSubdomain)
-	bb.WriteByte('.')
-	bb.Write(sdid)
-
-	qname := bb.Bytes()
-	dnsMsg, err := dnsClient.Lookup(dns.QueryTypeTXT, dns.QueryClassIN, qname)
+	dnsMsg, err := dnsClient.Lookup(dns.QueryTypeTXT, dns.QueryClassIN,
+		[]byte(dname))
 	if err != nil {
 		dnsClientPool.Put(dnsClient)
 		return nil, fmt.Errorf("dkim: LookupKey: " + err.Error())
 	}
 	if len(dnsMsg.Answer) == 0 {
 		dnsClientPool.Put(dnsClient)
-		return nil, fmt.Errorf("dkim: LookupKey: empty answer on '%s'", qname)
+		return nil, fmt.Errorf("dkim: LookupKey: empty answer on '%s'", dname)
 	}
 
 	dnsClientPool.Put(dnsClient)
 
 	answers := dnsMsg.FilterAnswers(dns.QueryTypeTXT)
 	if len(answers) == 0 {
-		return nil, fmt.Errorf("dkim: LookupKey: no TXT record on '%s'", qname)
+		return nil, fmt.Errorf("dkim: LookupKey: no TXT record on '%s'", dname)
 	}
 	if len(answers) != 1 {
-		return nil, fmt.Errorf("dkim: LookupKey: multiple TXT record on '%s'", qname)
+		return nil, fmt.Errorf("dkim: LookupKey: multiple TXT records on '%s'", dname)
 	}
 
 	txt := answers[0].RData().([]byte)
@@ -176,10 +167,10 @@ func newDNSClientPool() (err error) {
 }
 
 //
-// Bytes return text representation of Key.
+// Pack the key to be used in DNS TXT record.
 //
-func (key *Key) Bytes() []byte {
-	var bb bytes.Buffer
+func (key *Key) Pack() string {
+	var bb strings.Builder
 
 	if len(key.Version) > 0 {
 		bb.WriteString("v=")
@@ -219,7 +210,15 @@ func (key *Key) Bytes() []byte {
 		bb.Write(packKeyFlags(key.Flags))
 	}
 
-	return bb.Bytes()
+	return bb.String()
+}
+
+//
+// IsExpired will return true if key ExpiredAt time is less than current time;
+// otherwise it will return false.
+//
+func (key *Key) IsExpired() bool {
+	return key.ExpiredAt < time.Now().Unix()
 }
 
 func (key *Key) set(t *tag) (err error) {
diff --git a/lib/email/dkim/key_test.go b/lib/email/dkim/key_test.go
index 9e1817f0..e8069e6d 100644
--- a/lib/email/dkim/key_test.go
+++ b/lib/email/dkim/key_test.go
@@ -15,31 +15,31 @@ func TestLookupKey(t *testing.T) {
 	qmethod := QueryMethod{}
 
 	cases := []struct {
-		sdid     string
-		selector string
-		exp      string
-		expErr   string
+		domainName string
+		exp        string
+		expErr     string
 	}{{
-		expErr: "dkim: LookupKey: empty SDID",
+		expErr: "dkim: LookupKey: empty domain name",
 	}, {
-		sdid:     "amazonses.com",
-		selector: "ug7nbtf4gccmlpwj322ax3p6ow6yfsug",
-		exp:      "p=MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQCKkjP6XucgQ06cVZ89Ue/sQDu4v1/AJVd6mMK4bS2YmXk5PzWw4KWtWNUZlg77hegAChx1pG85lUbJ+x4awp28VXqRi3/jZoC6W+3ELysDvVohZPMRMadc+KVtyTiTH4BL38/8ZV9zkj4ZIaaYyiLAiYX+c3+lZQEF3rKDptRcpwIDAQAB; k=rsa;",
+		domainName: "ug7nbtf4gccmlpwj322ax3p6ow6yfsug._domainkey.amazonses.com",
+		exp:        "p=MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQCKkjP6XucgQ06cVZ89Ue/sQDu4v1/AJVd6mMK4bS2YmXk5PzWw4KWtWNUZlg77hegAChx1pG85lUbJ+x4awp28VXqRi3/jZoC6W+3ELysDvVohZPMRMadc+KVtyTiTH4BL38/8ZV9zkj4ZIaaYyiLAiYX+c3+lZQEF3rKDptRcpwIDAQAB; k=rsa;",
 	}, {
-		sdid:     "wikimedia-or-id.20150623.gappssmtp.com",
-		selector: "20150623",
-		exp:      "v=DKIM1; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA2UMfREvlgajdSp3jv1tJ9nLpi/mRYnGyKC3inEQ9a7zqUjLq/yXukgpXs9AEHlvBvioxlgAVCPQQsuc1xp9+KXQGgJ8jTsn5OtKm8u+YBCt6OfvpeCpvt0l9JXMMHBNYV4c0XiPE5RHX2ltI0Av20CfEy+vMecpFtVDg4rMngjLws/ro6qT63S20A4zyVs/V19WW5F2Lulgv+l+EJzz9XummIJHOlU5n5ChcWU3Rw5RVGTtNjTZnFUaNXly3fW0ahKcG5Qc3e0Rhztp57JJQTl3OmHiMR5cHsCnrl1VnBi3kaOoQBYsSuBm+KRhMIw/X9wkLY67VLdkrwlX3xxsp6wIDAQAB; k=rsa;",
+		domainName: "20150623._domainkey.wikimedia-or-id.20150623.gappssmtp.com",
+		exp:        "v=DKIM1; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA2UMfREvlgajdSp3jv1tJ9nLpi/mRYnGyKC3inEQ9a7zqUjLq/yXukgpXs9AEHlvBvioxlgAVCPQQsuc1xp9+KXQGgJ8jTsn5OtKm8u+YBCt6OfvpeCpvt0l9JXMMHBNYV4c0XiPE5RHX2ltI0Av20CfEy+vMecpFtVDg4rMngjLws/ro6qT63S20A4zyVs/V19WW5F2Lulgv+l+EJzz9XummIJHOlU5n5ChcWU3Rw5RVGTtNjTZnFUaNXly3fW0ahKcG5Qc3e0Rhztp57JJQTl3OmHiMR5cHsCnrl1VnBi3kaOoQBYsSuBm+KRhMIw/X9wkLY67VLdkrwlX3xxsp6wIDAQAB; k=rsa;",
 	}}
 
 	for _, c := range cases {
-		t.Logf("%s._domainkey.%s", c.selector, c.sdid)
+		t.Log(c.domainName)
 
-		got, err := LookupKey(qmethod, []byte(c.sdid), []byte(c.selector))
+		got, err := LookupKey(qmethod, c.domainName)
 		if err != nil {
 			test.Assert(t, "error", c.expErr, err.Error(), true)
 			continue
 		}
+		if got == nil {
+			continue
+		}
 
-		test.Assert(t, "Key", c.exp, string(got.Bytes()), true)
+		test.Assert(t, "Key", c.exp, got.Pack(), true)
 	}
 }
diff --git a/lib/email/dkim/keypool.go b/lib/email/dkim/keypool.go
new file mode 100644
index 00000000..eacd7f7a
--- /dev/null
+++ b/lib/email/dkim/keypool.go
@@ -0,0 +1,100 @@
+// Copyright 2019, Shulhan <ms@kilabit.info>. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package dkim
+
+import (
+	"sort"
+	"strconv"
+	"strings"
+	"sync"
+)
+
+//
+// KeyPool maintain cached DKIM public keys.
+//
+type KeyPool struct {
+	sync.Mutex
+	pool map[string]*Key
+}
+
+//
+// Clear the contents of key pool.
+//
+func (kp *KeyPool) Clear() {
+	kp.Lock()
+	for k := range kp.pool {
+		delete(kp.pool, k)
+	}
+	kp.Unlock()
+}
+
+//
+// Get cached DKIM key from pool or lookup using DNS/TXT method if not exist.
+//
+func (kp *KeyPool) Get(dname string) (key *Key, err error) {
+	if len(dname) == 0 {
+		return nil, nil
+	}
+
+	kp.Lock()
+	key, ok := kp.pool[dname]
+	if ok {
+		if !key.IsExpired() {
+			kp.Unlock()
+			return key, nil
+		}
+	}
+	kp.Unlock()
+
+	key, err = LookupKey(QueryMethod{}, dname)
+	if err != nil {
+		return nil, err
+	}
+
+	return key, nil
+}
+
+//
+// Put key to pool based on DKIM domain name ("d=" value plus "s=" value).
+//
+func (kp *KeyPool) Put(dname string, key *Key) {
+	if len(dname) == 0 || key == nil {
+		return
+	}
+	kp.Lock()
+	kp.pool[dname] = key
+	kp.Unlock()
+}
+
+//
+// String return text representation of DKIM key inside pool sorted by domain
+// name.  Each key is printed with the following format:
+// "{DomainName:ExpiredAt}"
+//
+func (kp *KeyPool) String() string {
+	var sb strings.Builder
+	kp.Lock()
+
+	dnames := make([]string, 0, len(kp.pool))
+	for k := range kp.pool {
+		dnames = append(dnames, k)
+	}
+
+	sort.Strings(dnames)
+
+	sb.WriteByte('[')
+	for _, v := range dnames {
+		key := kp.pool[v]
+		sb.WriteByte('{')
+		sb.WriteString(v)
+		sb.WriteByte(' ')
+		sb.WriteString(strconv.FormatInt(key.ExpiredAt, 10))
+		sb.WriteByte('}')
+	}
+	sb.WriteByte(']')
+
+	kp.Unlock()
+	return sb.String()
+}
diff --git a/lib/email/dkim/keypool_test.go b/lib/email/dkim/keypool_test.go
new file mode 100644
index 00000000..b5b828dd
--- /dev/null
+++ b/lib/email/dkim/keypool_test.go
@@ -0,0 +1,105 @@
+// Copyright 2019, Shulhan <ms@kilabit.info>. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package dkim
+
+import (
+	"testing"
+
+	"github.com/shuLhan/share/lib/test"
+)
+
+func TestKeyPoolClear(t *testing.T) {
+	DefaultKeyPool.Put("example.com", &Key{ExpiredAt: 1})
+	got := DefaultKeyPool.String()
+	test.Assert(t, "DefaultKeyPool.Clear", "[{example.com 1}]", got, true)
+
+	DefaultKeyPool.Clear()
+
+	got = DefaultKeyPool.String()
+	test.Assert(t, "DefaultKeyPool.Clear", "[]", got, true)
+}
+
+func TestKeyPoolPut(t *testing.T) {
+
+	cases := []struct {
+		dname string
+		key   *Key
+		exp   string
+	}{{
+		dname: "",
+		exp:   "[]",
+	}, {
+		dname: "emptykey",
+		exp:   "[]",
+	}, {
+		dname: "example.com",
+		key: &Key{
+			Public:    []byte("example.com"),
+			ExpiredAt: 1,
+		},
+		exp: "[{example.com 1}]",
+	}, {
+		dname: "example.com",
+		key: &Key{
+			Public:    []byte("example.com"),
+			ExpiredAt: 1577811600, // 2020-01-01
+		},
+		exp: "[{example.com 1577811600}]",
+	}, {
+		dname: "example.net",
+		key: &Key{
+			Public:    []byte("example.net"),
+			ExpiredAt: 1577811600, // 2020-01-01
+		},
+		exp: "[{example.com 1577811600}{example.net 1577811600}]",
+	}}
+
+	for _, c := range cases {
+		t.Log(c.dname)
+
+		DefaultKeyPool.Put(c.dname, c.key)
+		got := DefaultKeyPool.String()
+
+		test.Assert(t, "DefaultKeyPool", c.exp, got, true)
+	}
+}
+
+func TestKeyPoolGet(t *testing.T) {
+	cases := []struct {
+		dname  string
+		exp    string
+		expErr string
+	}{{
+		dname: "",
+	}, {
+		dname: "example.com",
+		exp:   "p=example.com; k=rsa;",
+	}, {
+		dname: "example.net",
+		exp:   "p=example.net; k=rsa;",
+	}, {
+		dname:  "example.org",
+		expErr: "dkim: LookupKey: multiple TXT records on 'example.org'",
+	}, {
+		dname: "ug7nbtf4gccmlpwj322ax3p6ow6yfsug._domainkey.amazonses.com",
+		exp:   "p=MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQCKkjP6XucgQ06cVZ89Ue/sQDu4v1/AJVd6mMK4bS2YmXk5PzWw4KWtWNUZlg77hegAChx1pG85lUbJ+x4awp28VXqRi3/jZoC6W+3ELysDvVohZPMRMadc+KVtyTiTH4BL38/8ZV9zkj4ZIaaYyiLAiYX+c3+lZQEF3rKDptRcpwIDAQAB; k=rsa;",
+	}}
+
+	for _, c := range cases {
+		t.Log(c.dname)
+
+		key, err := DefaultKeyPool.Get(c.dname)
+		if err != nil {
+			test.Assert(t, "error", c.expErr, err.Error(), true)
+			continue
+		}
+		if key == nil {
+			continue
+		}
+
+		got := key.Pack()
+		test.Assert(t, "DefaultKeyPool.Get", c.exp, got, true)
+	}
+}
diff --git a/lib/email/dkim/querymethod.go b/lib/email/dkim/querymethod.go
index b1d764fc..9722307b 100644
--- a/lib/email/dkim/querymethod.go
+++ b/lib/email/dkim/querymethod.go
@@ -5,7 +5,8 @@
 package dkim
 
 //
-// QueryMethod define a method to retrieve public key.
+// QueryMethod define a type and option to retrieve public key.
+// An empty QueryMethod will use default type and option, which is "dns/txt".
 //
 type QueryMethod struct {
 	Type   QueryType