add more tests for rest (#462)

core/discov/kubernetes/etcd-statefulset.yaml

@@ -0,0 +1,64 @@
+apiVersion: apps/v1
+kind: StatefulSet
+metadata:
+  name: "etcd"
+  namespace: discov
+  labels:
+    app: "etcd"
+spec:
+  serviceName: "etcd"
+  replicas: 5
+  template:
+    metadata:
+      name: "etcd"
+      labels:
+        app: "etcd"
+    spec:
+      volumes:
+      - name: etcd-pvc
+        persistentVolumeClaim:
+          claimName: etcd-pvc
+      containers:
+      - name: "etcd"
+        image: quay.io/coreos/etcd:latest
+        ports:
+        - containerPort: 2379
+          name: client
+        - containerPort: 2380
+          name: peer
+        env:
+        - name: CLUSTER_SIZE
+          value: "5"
+        - name: SET_NAME
+          value: "etcd"
+        - name: VOLNAME
+          valueFrom:
+            fieldRef:
+              apiVersion: v1
+              fieldPath: metadata.name
+        volumeMounts:
+        - name: etcd-pvc
+          mountPath: /var/lib/etcd
+          subPathExpr: $(VOLNAME)  # data mounted respectively in each pod
+        command:
+          - "/bin/sh"
+          - "-ecx"
+          - |
+
+            chmod 700 /var/lib/etcd
+
+            IP=$(hostname -i)
+            PEERS=""
+            for i in $(seq 0 $((${CLUSTER_SIZE} - 1))); do
+                PEERS="${PEERS}${PEERS:+,}${SET_NAME}-${i}=http://${SET_NAME}-${i}.${SET_NAME}:2380"
+            done
+            exec etcd --name ${HOSTNAME} \
+              --listen-peer-urls http://0.0.0.0:2380 \
+              --listen-client-urls http://0.0.0.0:2379 \
+              --advertise-client-urls http://${HOSTNAME}.${SET_NAME}.discov:2379 \
+              --initial-advertise-peer-urls http://${HOSTNAME}.${SET_NAME}:2380 \
+              --initial-cluster ${PEERS} \
+              --initial-cluster-state new \
+              --logger zap \
+              --data-dir /var/lib/etcd \
+              --auto-compaction-retention 1

rest/internal/security/contentsecurity_test.go

@@ -0,0 +1,169 @@
+package security
+
+import (
+	"crypto/hmac"
+	"crypto/md5"
+	"crypto/sha256"
+	"encoding/base64"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"os"
+	"strconv"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/codec"
+	"github.com/tal-tech/go-zero/core/fs"
+	"github.com/tal-tech/go-zero/rest/httpx"
+)
+
+const (
+	pubKey = `-----BEGIN PUBLIC KEY-----
+MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQCyeDYV2ieOtNDi6tuNtAbmUjN9
+pTHluAU5yiKEz8826QohcxqUKP3hybZBcm60p+rUxMAJFBJ8Dt+UJ6sEMzrf1rOF
+YOImVvORkXjpFU7sCJkhnLMs/kxtRzcZJG6ADUlG4GDCNcZpY/qELEvwgm2kCcHi
+tGC2mO8opFFFHTR0aQIDAQAB
+-----END PUBLIC KEY-----`
+	priKey = `-----BEGIN RSA PRIVATE KEY-----
+MIICXQIBAAKBgQCyeDYV2ieOtNDi6tuNtAbmUjN9pTHluAU5yiKEz8826QohcxqU
+KP3hybZBcm60p+rUxMAJFBJ8Dt+UJ6sEMzrf1rOFYOImVvORkXjpFU7sCJkhnLMs
+/kxtRzcZJG6ADUlG4GDCNcZpY/qELEvwgm2kCcHitGC2mO8opFFFHTR0aQIDAQAB
+AoGAcENv+jT9VyZkk6karLuG75DbtPiaN5+XIfAF4Ld76FWVOs9V88cJVON20xpx
+ixBphqexCMToj8MnXuHJEN5M9H15XXx/9IuiMm3FOw0i6o0+4V8XwHr47siT6T+r
+HuZEyXER/2qrm0nxyC17TXtd/+TtpfQWSbivl6xcAEo9RRECQQDj6OR6AbMQAIDn
+v+AhP/y7duDZimWJIuMwhigA1T2qDbtOoAEcjv3DB1dAswJ7clcnkxI9a6/0RDF9
+0IEHUcX9AkEAyHdcegWiayEnbatxWcNWm1/5jFnCN+GTRRFrOhBCyFr2ZdjFV4T+
+acGtG6omXWaZJy1GZz6pybOGy93NwLB93QJARKMJ0/iZDbOpHqI5hKn5mhd2Je25
+IHDCTQXKHF4cAQ+7njUvwIMLx2V5kIGYuMa5mrB/KMI6rmyvHv3hLewhnQJBAMMb
+cPUOENMllINnzk2oEd3tXiscnSvYL4aUeoErnGP2LERZ40/YD+mMZ9g6FVboaX04
+0oHf+k5mnXZD7WJyJD0CQQDJ2HyFbNaUUHK+lcifCibfzKTgmnNh9ZpePFumgJzI
+EfFE5H+nzsbbry2XgJbWzRNvuFTOLWn4zM+aFyy9WvbO
+-----END RSA PRIVATE KEY-----`
+	body = "hello world!"
+)
+
+var key = []byte("q4t7w!z%C*F-JaNdRgUjXn2r5u8x/A?D")
+
+func TestContentSecurity(t *testing.T) {
+	tests := []struct {
+		name        string
+		mode        string
+		extraKey    string
+		extraSecret string
+		extraTime   string
+		err         error
+		code        int
+	}{
+		{
+			name: "encrypted",
+			mode: "1",
+		},
+		{
+			name: "unencrypted",
+			mode: "0",
+		},
+		{
+			name: "bad content type",
+			mode: "a",
+			err:  ErrInvalidContentType,
+		},
+		{
+			name:        "bad secret",
+			mode:        "1",
+			extraSecret: "any",
+			err:         ErrInvalidSecret,
+		},
+		{
+			name:     "bad key",
+			mode:     "1",
+			extraKey: "any",
+			err:      ErrInvalidKey,
+		},
+		{
+			name:      "bad time",
+			mode:      "1",
+			extraTime: "any",
+			code:      httpx.CodeSignatureInvalidHeader,
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			r, err := http.NewRequest(http.MethodPost, "http://localhost:3333/a/b?c=first&d=second",
+				strings.NewReader(body))
+			assert.Nil(t, err)
+
+			timestamp := time.Now().Unix()
+			sha := sha256.New()
+			sha.Write([]byte(body))
+			bodySign := fmt.Sprintf("%x", sha.Sum(nil))
+			contentOfSign := strings.Join([]string{
+				strconv.FormatInt(timestamp, 10),
+				http.MethodPost,
+				r.URL.Path,
+				r.URL.RawQuery,
+				bodySign,
+			}, "\n")
+			sign := hs256(key, contentOfSign)
+			content := strings.Join([]string{
+				"version=v1",
+				"type=" + test.mode,
+				fmt.Sprintf("key=%s", base64.StdEncoding.EncodeToString(key)) + test.extraKey,
+				"time=" + strconv.FormatInt(timestamp, 10) + test.extraTime,
+			}, "; ")
+
+			encrypter, err := codec.NewRsaEncrypter([]byte(pubKey))
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			output, err := encrypter.Encrypt([]byte(content))
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			encryptedContent := base64.StdEncoding.EncodeToString(output)
+			r.Header.Set("X-Content-Security", strings.Join([]string{
+				fmt.Sprintf("key=%s", fingerprint(pubKey)),
+				"secret=" + encryptedContent + test.extraSecret,
+				"signature=" + sign,
+			}, "; "))
+
+			file, err := fs.TempFilenameWithText(priKey)
+			assert.Nil(t, err)
+			defer os.Remove(file)
+
+			dec, err := codec.NewRsaDecrypter(file)
+			assert.Nil(t, err)
+
+			header, err := ParseContentSecurity(map[string]codec.RsaDecrypter{
+				fingerprint(pubKey): dec,
+			}, r)
+			assert.Equal(t, test.err, err)
+			if err != nil {
+				return
+			}
+
+			assert.Equal(t, test.code, VerifySignature(r, header, time.Minute))
+		})
+	}
+}
+
+func fingerprint(key string) string {
+	h := md5.New()
+	io.WriteString(h, key)
+	return base64.StdEncoding.EncodeToString(h.Sum(nil))
+}
+
+func hs256(key []byte, body string) string {
+	h := hmac.New(sha256.New, key)
+	io.WriteString(h, body)
+	return base64.StdEncoding.EncodeToString(h.Sum(nil))
+}