You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
go-zero/core/httphandler/cryptionhandler_test.go

91 lines
2.5 KiB
Go

package httphandler
import (
"bytes"
"encoding/base64"
"io/ioutil"
"log"
"net/http"
"net/http/httptest"
"testing"
"zero/core/codec"
"github.com/stretchr/testify/assert"
)
const (
reqText = "ping"
respText = "pong"
)
var aesKey = []byte(`PdSgVkYp3s6v9y$B&E)H+MbQeThWmZq4`)
func init() {
log.SetOutput(ioutil.Discard)
}
func TestCryptionHandlerGet(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "/any", nil)
handler := CryptionHandler(aesKey)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, err := w.Write([]byte(respText))
w.Header().Set("X-Test", "test")
assert.Nil(t, err)
}))
recorder := httptest.NewRecorder()
handler.ServeHTTP(recorder, req)
expect, err := codec.EcbEncrypt(aesKey, []byte(respText))
assert.Nil(t, err)
assert.Equal(t, http.StatusOK, recorder.Code)
assert.Equal(t, "test", recorder.Header().Get("X-Test"))
assert.Equal(t, base64.StdEncoding.EncodeToString(expect), recorder.Body.String())
}
func TestCryptionHandlerPost(t *testing.T) {
var buf bytes.Buffer
enc, err := codec.EcbEncrypt(aesKey, []byte(reqText))
assert.Nil(t, err)
buf.WriteString(base64.StdEncoding.EncodeToString(enc))
req := httptest.NewRequest(http.MethodPost, "/any", &buf)
handler := CryptionHandler(aesKey)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body, err := ioutil.ReadAll(r.Body)
assert.Nil(t, err)
assert.Equal(t, reqText, string(body))
w.Write([]byte(respText))
}))
recorder := httptest.NewRecorder()
handler.ServeHTTP(recorder, req)
expect, err := codec.EcbEncrypt(aesKey, []byte(respText))
assert.Nil(t, err)
assert.Equal(t, http.StatusOK, recorder.Code)
assert.Equal(t, base64.StdEncoding.EncodeToString(expect), recorder.Body.String())
}
func TestCryptionHandlerPostBadEncryption(t *testing.T) {
var buf bytes.Buffer
enc, err := codec.EcbEncrypt(aesKey, []byte(reqText))
assert.Nil(t, err)
buf.Write(enc)
req := httptest.NewRequest(http.MethodPost, "/any", &buf)
handler := CryptionHandler(aesKey)(nil)
recorder := httptest.NewRecorder()
handler.ServeHTTP(recorder, req)
assert.Equal(t, http.StatusBadRequest, recorder.Code)
}
func TestCryptionHandlerWriteHeader(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "/any", nil)
handler := CryptionHandler(aesKey)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusServiceUnavailable)
}))
recorder := httptest.NewRecorder()
handler.ServeHTTP(recorder, req)
assert.Equal(t, http.StatusServiceUnavailable, recorder.Code)
}