From 4211672bfd3f3a4eeb9325045dadf38fc076c562 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Wed, 20 Sep 2023 00:01:26 +0800
Subject: [PATCH] chore: add more tests (#3577)

---
 rest/handler/cryptionhandler.go      |  4 +--
 rest/handler/cryptionhandler_test.go | 37 ++++++++++++++++++++++++++--
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/rest/handler/cryptionhandler.go b/rest/handler/cryptionhandler.go
index 0ef92379..18ba85b3 100644
--- a/rest/handler/cryptionhandler.go
+++ b/rest/handler/cryptionhandler.go
@@ -132,7 +132,7 @@ func (w *cryptionResponseWriter) flush(key []byte) {
 	body := base64.StdEncoding.EncodeToString(content)
 	if n, err := io.WriteString(w.ResponseWriter, body); err != nil {
 		logx.Errorf("write response failed, error: %s", err)
-	} else if n < len(content) {
-		logx.Errorf("actual bytes: %d, written bytes: %d", len(content), n)
+	} else if n < len(body) {
+		logx.Errorf("actual bytes: %d, written bytes: %d", len(body), n)
 	}
 }
diff --git a/rest/handler/cryptionhandler_test.go b/rest/handler/cryptionhandler_test.go
index b24a5850..77723126 100644
--- a/rest/handler/cryptionhandler_test.go
+++ b/rest/handler/cryptionhandler_test.go
@@ -2,12 +2,13 @@ package handler
 
 import (
 	"bytes"
+	"crypto/rand"
 	"encoding/base64"
 	"io"
-	"math/rand"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"testing/iotest"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/codec"
@@ -37,6 +38,19 @@ func TestCryptionHandlerGet(t *testing.T) {
 	assert.Equal(t, base64.StdEncoding.EncodeToString(expect), recorder.Body.String())
 }
 
+func TestCryptionHandlerGet_badKey(t *testing.T) {
+	req := httptest.NewRequest(http.MethodGet, "/any", http.NoBody)
+	handler := CryptionHandler(append(aesKey, 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)
+	assert.Equal(t, http.StatusInternalServerError, recorder.Code)
+}
+
 func TestCryptionHandlerPost(t *testing.T) {
 	var buf bytes.Buffer
 	enc, err := codec.EcbEncrypt(aesKey, []byte(reqText))
@@ -120,10 +134,29 @@ func TestCryptionHandler_ContentTooLong(t *testing.T) {
 	defer svr.Close()
 
 	body := make([]byte, maxBytes+1)
-	rand.Read(body)
+	_, err := rand.Read(body)
+	assert.NoError(t, err)
 	req, err := http.NewRequest(http.MethodPost, svr.URL, bytes.NewReader(body))
 	assert.Nil(t, err)
 	resp, err := http.DefaultClient.Do(req)
 	assert.Nil(t, err)
 	assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
 }
+
+func TestCryptionHandler_BadBody(t *testing.T) {
+	req, err := http.NewRequest(http.MethodPost, "/foo", iotest.ErrReader(io.ErrUnexpectedEOF))
+	assert.Nil(t, err)
+	err = decryptBody(maxBytes, aesKey, req)
+	assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
+}
+
+func TestCryptionHandler_BadKey(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)
+	err = decryptBody(maxBytes, append(aesKey, aesKey...), req)
+	assert.Error(t, err)
+}