From 7b39e101bf383d186e40479587b4f0e03df277c2 Mon Sep 17 00:00:00 2001
From: jkskj <lllljkskj@163.com>
Date: Wed, 10 Jul 2024 12:40:09 +0800
Subject: [PATCH] fix: test error and decompressMiddleware doesn't work (#20)

* fix: test error and decompressMiddleware doesn't work

* fix: return error in ClientMiddleware

* chore: optimize license header
---
 client_middleware.go     | 10 +++++++---
 example/standard/main.go |  2 +-
 example/stream/main.go   |  2 +-
 gzip.go                  |  2 +-
 gzip_test.go             |  6 ++++--
 options.go               |  8 +++++---
 srv_middleware.go        |  2 +-
 srv_stream_middleware.go |  2 +-
 8 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/client_middleware.go b/client_middleware.go
index ab00e78..eb53a2a 100644
--- a/client_middleware.go
+++ b/client_middleware.go
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 CloudWeGo Authors
+ * Copyright 2024 CloudWeGo Authors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -82,10 +82,14 @@ func (g *gzipClientMiddleware) ClientMiddleware(next client.Endpoint) client.End
 
 		err = next(ctx, req, resp)
 		if err != nil {
-			return
+			return err
 		}
 		if fn := g.DecompressFnForClient; fn != nil && strings.EqualFold(resp.Header.Get("Content-Encoding"), "gzip") {
-			fn(next)
+			f := fn(next)
+			err = f(ctx, req, resp)
+			if err != nil {
+				return err
+			}
 		}
 		return nil
 	}
diff --git a/example/standard/main.go b/example/standard/main.go
index d3a3d17..b30ab8d 100644
--- a/example/standard/main.go
+++ b/example/standard/main.go
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 CloudWeGo Authors
+ * Copyright 2024 CloudWeGo Authors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/example/stream/main.go b/example/stream/main.go
index 80eedc3..336badc 100644
--- a/example/stream/main.go
+++ b/example/stream/main.go
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 CloudWeGo Authors
+ * Copyright 2024 CloudWeGo Authors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/gzip.go b/gzip.go
index 5c47ad4..3f25f1c 100644
--- a/gzip.go
+++ b/gzip.go
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 CloudWeGo Authors
+ * Copyright 2024 CloudWeGo Authors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/gzip_test.go b/gzip_test.go
index 99610ca..c0468a5 100644
--- a/gzip_test.go
+++ b/gzip_test.go
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 CloudWeGo Authors
+ * Copyright 2024 CloudWeGo Authors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -405,11 +405,12 @@ func TestNoGzipForClient(t *testing.T) {
 func TestDecompressGzipForClient(t *testing.T) {
 	h := server.Default(server.WithHostPorts("127.0.0.1:2338"))
 
+	h.Use(Gzip(DefaultCompression, WithDecompressFn(DefaultDecompressHandle)))
 	h.GET("/", func(ctx context.Context, c *app.RequestContext) {
 		c.Header("Content-Length", strconv.Itoa(len(testResponse)))
 		c.String(200, testResponse)
 	})
-	h.Use(Gzip(DefaultCompression, WithDecompressFn(DefaultDecompressHandle)))
+
 	go h.Spin()
 
 	time.Sleep(time.Second)
@@ -425,6 +426,7 @@ func TestDecompressGzipForClient(t *testing.T) {
 
 	req.SetBodyString("bar")
 	req.SetRequestURI("http://127.0.0.1:2338/")
+	req.SetHeader("Accept-Encoding", "gzip")
 
 	err = cli.Do(context.Background(), req, res)
 	if err != nil {
diff --git a/options.go b/options.go
index 80e3689..89f5489 100644
--- a/options.go
+++ b/options.go
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 CloudWeGo Authors
+ * Copyright 2024 CloudWeGo Authors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -41,6 +41,7 @@
 package gzip
 
 import (
+	"bytes"
 	"context"
 	"net/http"
 	"regexp"
@@ -216,7 +217,8 @@ func DefaultDecompressMiddlewareForClient(next client.Endpoint) client.Endpoint
 		}
 		resp.Header.DelBytes([]byte("Content-Encoding"))
 		resp.Header.DelBytes([]byte("Content-Length"))
-		resp.SetBody(gunzipBytes)
-		return next(ctx, req, resp)
+		resp.Header.DelBytes([]byte("Vary"))
+		resp.SetBodyStream(bytes.NewBuffer(gunzipBytes), len(gunzipBytes))
+		return nil
 	}
 }
diff --git a/srv_middleware.go b/srv_middleware.go
index 8807c87..4664933 100644
--- a/srv_middleware.go
+++ b/srv_middleware.go
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 CloudWeGo Authors
+ * Copyright 2024 CloudWeGo Authors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/srv_stream_middleware.go b/srv_stream_middleware.go
index 682ce49..fa87552 100644
--- a/srv_stream_middleware.go
+++ b/srv_stream_middleware.go
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 CloudWeGo Authors
+ * Copyright 2024 CloudWeGo Authors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.