From 91c8eb4e9fecbc7d4b3c944ad212ed68e76fc992 Mon Sep 17 00:00:00 2001
From: moocss <moocss@gmail.com>
Date: Sun, 1 Oct 2023 13:01:34 +0800
Subject: [PATCH] add: transport ws

---
 internal/testdata/encoding/test.pb.go         |  5 +-
 internal/testdata/helloworld/helloworld.pb.go |  5 +-
 .../testdata/helloworld/helloworld_grpc.pb.go |  1 +
 pkg/utils/dir_test.go                         |  3 +-
 plugins/transport/websocket/client.go         | 13 ++-
 plugins/transport/websocket/encoding.go       | 49 +++++++++++
 plugins/transport/websocket/go.mod            | 34 ++------
 plugins/transport/websocket/go.sum            | 39 +++++++++
 plugins/transport/websocket/logger.go         |  2 +-
 plugins/transport/websocket/options.go        | 56 ++++++++++++-
 plugins/transport/websocket/server.go         | 84 +++----------------
 plugins/transport/websocket/transport.go      | 14 +---
 plugins/transport/websocket/transport_test.go | 25 +-----
 transport/grpc/interceptor.go                 |  2 +-
 transport/grpc/server.go                      | 12 +--
 transport/http/server.go                      |  5 --
 transport/transport.go                        |  1 -
 17 files changed, 186 insertions(+), 164 deletions(-)
 create mode 100644 plugins/transport/websocket/go.sum

diff --git a/internal/testdata/encoding/test.pb.go b/internal/testdata/encoding/test.pb.go
index b5b6c5d..a94093b 100644
--- a/internal/testdata/encoding/test.pb.go
+++ b/internal/testdata/encoding/test.pb.go
@@ -7,11 +7,12 @@
 package encoding
 
 import (
+	reflect "reflect"
+	sync "sync"
+
 	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
 	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
 	structpb "google.golang.org/protobuf/types/known/structpb"
-	reflect "reflect"
-	sync "sync"
 )
 
 const (
diff --git a/internal/testdata/helloworld/helloworld.pb.go b/internal/testdata/helloworld/helloworld.pb.go
index f807f09..709e316 100644
--- a/internal/testdata/helloworld/helloworld.pb.go
+++ b/internal/testdata/helloworld/helloworld.pb.go
@@ -7,11 +7,12 @@
 package helloworld
 
 import (
+	reflect "reflect"
+	sync "sync"
+
 	_ "google.golang.org/genproto/googleapis/api/annotations"
 	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
 	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
-	reflect "reflect"
-	sync "sync"
 )
 
 const (
diff --git a/internal/testdata/helloworld/helloworld_grpc.pb.go b/internal/testdata/helloworld/helloworld_grpc.pb.go
index 44b0382..68912bb 100644
--- a/internal/testdata/helloworld/helloworld_grpc.pb.go
+++ b/internal/testdata/helloworld/helloworld_grpc.pb.go
@@ -8,6 +8,7 @@ package helloworld
 
 import (
 	context "context"
+
 	grpc "google.golang.org/grpc"
 	codes "google.golang.org/grpc/codes"
 	status "google.golang.org/grpc/status"
diff --git a/pkg/utils/dir_test.go b/pkg/utils/dir_test.go
index 9e1532e..764f0e4 100644
--- a/pkg/utils/dir_test.go
+++ b/pkg/utils/dir_test.go
@@ -1,12 +1,13 @@
 package utils
 
 import (
-	"bou.ke/monkey"
 	"fmt"
 	"os"
 	"path/filepath"
 	"testing"
 
+	"bou.ke/monkey"
+
 	"github.com/stretchr/testify/assert"
 )
 
diff --git a/plugins/transport/websocket/client.go b/plugins/transport/websocket/client.go
index 3e4f2a3..958fe3e 100644
--- a/plugins/transport/websocket/client.go
+++ b/plugins/transport/websocket/client.go
@@ -6,11 +6,8 @@ import (
 	"net/url"
 	"time"
 
-	"github.com/go-kratos/kratos/v2/encoding"
-
+	"github.com/apus-run/sea-kit/encoding"
 	ws "github.com/gorilla/websocket"
-
-	"github.com/tx7do/kratos-transport/broker"
 )
 
 type ClientMessageHandler func(MessagePayload) error
@@ -123,7 +120,7 @@ func (c *Client) marshalMessage(messageType MessageType, message MessagePayload)
 	case PayloadTypeBinary:
 		var msg BinaryMessage
 		msg.Type = messageType
-		msg.Body, err = broker.Marshal(c.codec, message)
+		msg.Body, err = Marshal(c.codec, message)
 		if err != nil {
 			return nil, err
 		}
@@ -137,7 +134,7 @@ func (c *Client) marshalMessage(messageType MessageType, message MessagePayload)
 		var buf []byte
 		var msg TextMessage
 		msg.Type = messageType
-		buf, err = broker.Marshal(c.codec, message)
+		buf, err = Marshal(c.codec, message)
 		msg.Body = string(buf)
 		if err != nil {
 			return nil, err
@@ -257,7 +254,7 @@ func (c *Client) unmarshalMessage(buf []byte) (*ClientHandlerData, MessagePayloa
 			payload = msg.Body
 		}
 
-		if err := broker.Unmarshal(c.codec, msg.Body, &payload); err != nil {
+		if err := Unmarshal(c.codec, msg.Body, &payload); err != nil {
 			LogErrorf("unmarshal message exception: %s", err)
 			return nil, nil, err
 		}
@@ -283,7 +280,7 @@ func (c *Client) unmarshalMessage(buf []byte) (*ClientHandlerData, MessagePayloa
 			payload = msg.Body
 		}
 
-		if err := broker.Unmarshal(c.codec, []byte(msg.Body), &payload); err != nil {
+		if err := Unmarshal(c.codec, []byte(msg.Body), &payload); err != nil {
 			LogErrorf("unmarshal message exception: %s", err)
 			return nil, nil, err
 		}
diff --git a/plugins/transport/websocket/encoding.go b/plugins/transport/websocket/encoding.go
index 708bc8c..2e84152 100644
--- a/plugins/transport/websocket/encoding.go
+++ b/plugins/transport/websocket/encoding.go
@@ -1 +1,50 @@
 package websocket
+
+import (
+	"bytes"
+	"encoding/gob"
+	"errors"
+
+	"github.com/apus-run/sea-kit/encoding"
+	_ "github.com/apus-run/sea-kit/encoding/json"
+	_ "github.com/apus-run/sea-kit/encoding/proto"
+)
+
+func Marshal(codec encoding.Codec, msg Any) ([]byte, error) {
+	if msg == nil {
+		return nil, errors.New("message is nil")
+	}
+
+	if codec != nil {
+		dataBuffer, err := codec.Marshal(msg)
+		if err != nil {
+			return nil, err
+		}
+		return dataBuffer, nil
+	} else {
+		switch t := msg.(type) {
+		case []byte:
+			return t, nil
+		case string:
+			return []byte(t), nil
+		default:
+			var buf bytes.Buffer
+			enc := gob.NewEncoder(&buf)
+			if err := enc.Encode(msg); err != nil {
+				return nil, err
+			}
+			return buf.Bytes(), nil
+		}
+	}
+}
+
+func Unmarshal(codec encoding.Codec, inputData []byte, outValue interface{}) error {
+	if codec != nil {
+		if err := codec.Unmarshal(inputData, outValue); err != nil {
+			return err
+		}
+	} else if outValue == nil {
+		outValue = inputData
+	}
+	return nil
+}
diff --git a/plugins/transport/websocket/go.mod b/plugins/transport/websocket/go.mod
index 9e7f73b..118790b 100644
--- a/plugins/transport/websocket/go.mod
+++ b/plugins/transport/websocket/go.mod
@@ -1,42 +1,24 @@
-module github.com/tx7do/kratos-transport/transport/websocket
+module github.com/apus-run/gaia/plugins/transport/websocket
 
 go 1.19
 
 require (
-	github.com/go-kratos/kratos/v2 v2.7.0
+	github.com/apus-run/gaia v1.9.0
+	github.com/apus-run/sea-kit/encoding v0.0.0-20230930061415-4e76dbc0e8a9
+	github.com/apus-run/sea-kit/log v0.0.0-20230930061415-4e76dbc0e8a9
 	github.com/google/uuid v1.3.1
 	github.com/gorilla/websocket v1.5.0
 	github.com/stretchr/testify v1.8.4
-	github.com/tx7do/kratos-transport v1.0.12
 )
 
 require (
-	github.com/cenkalti/backoff/v4 v4.2.1 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/go-logr/logr v1.2.4 // indirect
-	github.com/go-logr/stdr v1.2.2 // indirect
-	github.com/go-playground/form/v4 v4.2.1 // indirect
-	github.com/golang/protobuf v1.5.3 // indirect
-	github.com/grpc-ecosystem/grpc-gateway/v2 v2.17.1 // indirect
-	github.com/openzipkin/zipkin-go v0.4.2 // indirect
+	github.com/kr/pretty v0.3.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	go.opentelemetry.io/otel v1.17.0 // indirect
-	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.17.0 // indirect
-	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.17.0 // indirect
-	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.17.0 // indirect
-	go.opentelemetry.io/otel/exporters/zipkin v1.17.0 // indirect
-	go.opentelemetry.io/otel/metric v1.17.0 // indirect
-	go.opentelemetry.io/otel/sdk v1.17.0 // indirect
-	go.opentelemetry.io/otel/trace v1.17.0 // indirect
-	go.opentelemetry.io/proto/otlp v1.0.0 // indirect
-	golang.org/x/net v0.15.0 // indirect
-	golang.org/x/sys v0.12.0 // indirect
-	golang.org/x/text v0.13.0 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20230822172742-b8732ec3820d // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20230822172742-b8732ec3820d // indirect
-	google.golang.org/grpc v1.58.0 // indirect
+	github.com/rogpeppe/go-internal v1.10.0 // indirect
 	google.golang.org/protobuf v1.31.0 // indirect
+	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 
-replace github.com/tx7do/kratos-transport => ../../
+replace github.com/apus-run/gaia => ../../../
diff --git a/plugins/transport/websocket/go.sum b/plugins/transport/websocket/go.sum
new file mode 100644
index 0000000..1b657fa
--- /dev/null
+++ b/plugins/transport/websocket/go.sum
@@ -0,0 +1,39 @@
+github.com/apus-run/sea-kit/encoding v0.0.0-20230930061415-4e76dbc0e8a9 h1:rsDikhV1ckW+GxFi5Vw9wG10rInlnS45TKK29oFF/g8=
+github.com/apus-run/sea-kit/encoding v0.0.0-20230930061415-4e76dbc0e8a9/go.mod h1:IUsjPu4FzcbbHR2exjACyTa4XCxQr8b9JdVkzM4r3xQ=
+github.com/apus-run/sea-kit/log v0.0.0-20230930061415-4e76dbc0e8a9 h1:fQLt1+Z1ifMgvONnkT7kFBQZcUttc8YqYBr+SjrqUXE=
+github.com/apus-run/sea-kit/log v0.0.0-20230930061415-4e76dbc0e8a9/go.mod h1:bkjkCOCQbbVy8HJbZ8HpVZ8yR36L9esmhEu869idCc8=
+github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
+github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
+github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/google/uuid v1.3.1 h1:KjJaJ9iWZ3jOFZIf1Lqf4laDRCasjl0BCmnEGxkdLb4=
+github.com/google/uuid v1.3.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
+github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
+github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
+github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
+github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
+github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
+github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
+github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
+github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
+github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ=
+github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog=
+github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
+golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
+google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8=
+google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
diff --git a/plugins/transport/websocket/logger.go b/plugins/transport/websocket/logger.go
index 79b5291..656e1be 100644
--- a/plugins/transport/websocket/logger.go
+++ b/plugins/transport/websocket/logger.go
@@ -3,7 +3,7 @@ package websocket
 import (
 	"fmt"
 
-	"github.com/go-kratos/kratos/v2/log"
+	"github.com/apus-run/sea-kit/log"
 )
 
 const (
diff --git a/plugins/transport/websocket/options.go b/plugins/transport/websocket/options.go
index 9b025cd..78b6ff8 100644
--- a/plugins/transport/websocket/options.go
+++ b/plugins/transport/websocket/options.go
@@ -3,9 +3,11 @@ package websocket
 import (
 	"crypto/tls"
 	"net"
+	"net/http"
 	"time"
 
-	"github.com/go-kratos/kratos/v2/encoding"
+	"github.com/apus-run/sea-kit/encoding"
+	ws "github.com/gorilla/websocket"
 )
 
 type PayloadType uint8
@@ -15,6 +17,58 @@ const (
 	PayloadTypeText   = 1
 )
 
+type Server struct {
+	*http.Server
+
+	lis      net.Listener
+	tlsConf  *tls.Config
+	upgrader *ws.Upgrader
+
+	network     string
+	address     string
+	path        string
+	strictSlash bool
+
+	timeout time.Duration
+
+	err   error
+	codec encoding.Codec
+
+	messageHandlers MessageHandlerMap
+
+	sessionMgr *SessionManager
+
+	register   chan *Session
+	unregister chan *Session
+
+	payloadType PayloadType
+}
+
+// defaultServer return a default config server
+func defaultServer() *Server {
+	return &Server{
+		network:     "tcp",
+		address:     ":0",
+		timeout:     1 * time.Second,
+		strictSlash: true,
+		path:        "/",
+
+		messageHandlers: make(MessageHandlerMap),
+
+		sessionMgr: NewSessionManager(),
+		upgrader: &ws.Upgrader{
+			ReadBufferSize:  1024,
+			WriteBufferSize: 1024,
+			CheckOrigin:     func(r *http.Request) bool { return true },
+		},
+
+		register:   make(chan *Session),
+		unregister: make(chan *Session),
+
+		payloadType: PayloadTypeBinary,
+	}
+}
+
 type ServerOption func(o *Server)
 
 func WithNetwork(network string) ServerOption {
diff --git a/plugins/transport/websocket/server.go b/plugins/transport/websocket/server.go
index 805f421..92030df 100644
--- a/plugins/transport/websocket/server.go
+++ b/plugins/transport/websocket/server.go
@@ -2,21 +2,14 @@ package websocket
 
 import (
 	"context"
-	"crypto/tls"
 	"encoding/json"
 	"errors"
 	"net"
 	"net/http"
 	"net/url"
 	"strings"
-	"time"
 
-	"github.com/go-kratos/kratos/v2/encoding"
-	"github.com/go-kratos/kratos/v2/transport"
-
-	ws "github.com/gorilla/websocket"
-
-	"github.com/tx7do/kratos-transport/broker"
+	"github.com/apus-run/gaia/transport"
 )
 
 type Binder func() Any
@@ -36,57 +29,18 @@ var (
 	_ transport.Endpointer = (*Server)(nil)
 )
 
-type Server struct {
-	*http.Server
-
-	lis      net.Listener
-	tlsConf  *tls.Config
-	upgrader *ws.Upgrader
-
-	network     string
-	address     string
-	path        string
-	strictSlash bool
-
-	timeout time.Duration
-
-	err   error
-	codec encoding.Codec
-
-	messageHandlers MessageHandlerMap
-
-	sessionMgr *SessionManager
-
-	register   chan *Session
-	unregister chan *Session
-
-	payloadType PayloadType
-}
-
 func NewServer(opts ...ServerOption) *Server {
-	srv := &Server{
-		network:     "tcp",
-		address:     ":0",
-		timeout:     1 * time.Second,
-		strictSlash: true,
-		path:        "/",
-
-		messageHandlers: make(MessageHandlerMap),
-
-		sessionMgr: NewSessionManager(),
-		upgrader: &ws.Upgrader{
-			ReadBufferSize:  1024,
-			WriteBufferSize: 1024,
-			CheckOrigin:     func(r *http.Request) bool { return true },
-		},
-
-		register:   make(chan *Session),
-		unregister: make(chan *Session),
+	srv := defaultServer()
+	// apply options
+	for _, o := range opts {
+		o(srv)
+	}
 
-		payloadType: PayloadTypeBinary,
+	srv.Server = &http.Server{
+		TLSConfig: srv.tlsConf,
 	}
 
-	srv.init(opts...)
+	http.HandleFunc(srv.path, srv.wsHandler)
 
 	srv.err = srv.listen()
 
@@ -97,18 +51,6 @@ func (s *Server) Name() string {
 	return string(KindWebsocket)
 }
 
-func (s *Server) init(opts ...ServerOption) {
-	for _, o := range opts {
-		o(s)
-	}
-
-	s.Server = &http.Server{
-		TLSConfig: s.tlsConf,
-	}
-
-	http.HandleFunc(s.path, s.wsHandler)
-}
-
 func (s *Server) SessionCount() int {
 	return s.sessionMgr.Count()
 }
@@ -153,7 +95,7 @@ func (s *Server) marshalMessage(messageType MessageType, message MessagePayload)
 	case PayloadTypeBinary:
 		var msg BinaryMessage
 		msg.Type = messageType
-		msg.Body, err = broker.Marshal(s.codec, message)
+		msg.Body, err = Marshal(s.codec, message)
 		if err != nil {
 			return nil, err
 		}
@@ -167,7 +109,7 @@ func (s *Server) marshalMessage(messageType MessageType, message MessagePayload)
 		var buf []byte
 		var msg TextMessage
 		msg.Type = messageType
-		buf, err = broker.Marshal(s.codec, message)
+		buf, err = Marshal(s.codec, message)
 		msg.Body = string(buf)
 		if err != nil {
 			return nil, err
@@ -252,7 +194,7 @@ func (s *Server) unmarshalMessage(buf []byte) (*HandlerData, MessagePayload, err
 			payload = msg.Body
 		}
 
-		if err := broker.Unmarshal(s.codec, msg.Body, &payload); err != nil {
+		if err := Unmarshal(s.codec, msg.Body, &payload); err != nil {
 			LogErrorf("unmarshal message exception: %s", err)
 			return nil, nil, err
 		}
@@ -278,7 +220,7 @@ func (s *Server) unmarshalMessage(buf []byte) (*HandlerData, MessagePayload, err
 			payload = msg.Body
 		}
 
-		if err := broker.Unmarshal(s.codec, []byte(msg.Body), &payload); err != nil {
+		if err := Unmarshal(s.codec, []byte(msg.Body), &payload); err != nil {
 			LogErrorf("unmarshal message exception: %s", err)
 			return nil, nil, err
 		}
diff --git a/plugins/transport/websocket/transport.go b/plugins/transport/websocket/transport.go
index 6cada46..2785866 100644
--- a/plugins/transport/websocket/transport.go
+++ b/plugins/transport/websocket/transport.go
@@ -4,7 +4,7 @@ import (
 	"context"
 	"net/http"
 
-	"github.com/go-kratos/kratos/v2/transport"
+	"github.com/apus-run/gaia/transport"
 )
 
 const (
@@ -23,8 +23,6 @@ type Transporter interface {
 type Transport struct {
 	endpoint     string
 	operation    string
-	reqHeader    headerCarrier
-	replyHeader  headerCarrier
 	request      *http.Request
 	pathTemplate string
 }
@@ -49,16 +47,6 @@ func (tr *Transport) Request() *http.Request {
 	return tr.request
 }
 
-// RequestHeader returns the request header.
-func (tr *Transport) RequestHeader() transport.Header {
-	return tr.reqHeader
-}
-
-// ReplyHeader returns the reply header.
-func (tr *Transport) ReplyHeader() transport.Header {
-	return tr.replyHeader
-}
-
 // PathTemplate returns the http path template.
 func (tr *Transport) PathTemplate() string {
 	return tr.pathTemplate
diff --git a/plugins/transport/websocket/transport_test.go b/plugins/transport/websocket/transport_test.go
index fc09403..3b9e225 100644
--- a/plugins/transport/websocket/transport_test.go
+++ b/plugins/transport/websocket/transport_test.go
@@ -2,11 +2,12 @@ package websocket
 
 import (
 	"context"
-	"github.com/go-kratos/kratos/v2/transport"
 	"net/http"
 	"reflect"
 	"sort"
 	"testing"
+
+	"github.com/apus-run/gaia/transport"
 )
 
 func TestTransport_Kind(t *testing.T) {
@@ -40,24 +41,6 @@ func TestTransport_Request(t *testing.T) {
 	}
 }
 
-func TestTransport_RequestHeader(t *testing.T) {
-	v := headerCarrier{}
-	v.Set("a", "1")
-	o := &Transport{reqHeader: v}
-	if !reflect.DeepEqual("1", o.RequestHeader().Get("a")) {
-		t.Errorf("expect %v, got %v", "1", o.RequestHeader().Get("a"))
-	}
-}
-
-func TestTransport_ReplyHeader(t *testing.T) {
-	v := headerCarrier{}
-	v.Set("a", "1")
-	o := &Transport{replyHeader: v}
-	if !reflect.DeepEqual("1", o.ReplyHeader().Get("a")) {
-		t.Errorf("expect %v, got %v", "1", o.ReplyHeader().Get("a"))
-	}
-}
-
 func TestTransport_PathTemplate(t *testing.T) {
 	v := "template"
 	o := &Transport{pathTemplate: v}
@@ -86,8 +69,8 @@ func TestHeaderCarrier_Keys(t *testing.T) {
 func TestSetOperation(t *testing.T) {
 	tr := &Transport{}
 	ctx := transport.NewServerContext(context.Background(), tr)
-	SetOperation(ctx, "kratos")
-	if !reflect.DeepEqual(tr.operation, "kratos") {
+	SetOperation(ctx, "gaia")
+	if !reflect.DeepEqual(tr.operation, "gaia") {
 		t.Errorf("expect %v, got %v", "kratos", tr.operation)
 	}
 }
diff --git a/transport/grpc/interceptor.go b/transport/grpc/interceptor.go
index 22b0445..606eb57 100644
--- a/transport/grpc/interceptor.go
+++ b/transport/grpc/interceptor.go
@@ -2,13 +2,13 @@ package grpc
 
 import (
 	"context"
-	"github.com/apus-run/gaia/transport"
 	"time"
 
 	"google.golang.org/grpc"
 
 	ic "github.com/apus-run/gaia/internal/context"
 	"github.com/apus-run/gaia/middleware"
+	"github.com/apus-run/gaia/transport"
 )
 
 // wrappedStream is rewrite grpc stream's context
diff --git a/transport/grpc/server.go b/transport/grpc/server.go
index 0708932..c71fcc8 100644
--- a/transport/grpc/server.go
+++ b/transport/grpc/server.go
@@ -109,21 +109,11 @@ func (s *Server) Stop(ctx context.Context) error {
 		s.adminClean()
 	}
 	s.health.Shutdown()
-	s.Server.Stop()
+	s.GracefulStop()
 	log.Info("[gRPC] server stopping")
 	return nil
 }
 
-func (s *Server) GracefullyStop(ctx context.Context) error {
-	if s.adminClean != nil {
-		s.adminClean()
-	}
-	s.health.Shutdown()
-	s.Server.GracefulStop()
-	log.Info("[gRPC] server graceful stopping")
-	return nil
-}
-
 func (s *Server) listenAndEndpoint() error {
 	if s.lis == nil {
 		lis, err := net.Listen(s.network, s.address)
diff --git a/transport/http/server.go b/transport/http/server.go
index 1ac69ba..483502e 100644
--- a/transport/http/server.go
+++ b/transport/http/server.go
@@ -104,11 +104,6 @@ func (s *Server) Start(ctx context.Context) error {
 
 // Stop stop the HTTP server.
 func (s *Server) Stop(ctx context.Context) error {
-	log.Infof("[HTTP] server is stopping")
-	return s.Close()
-}
-
-func (s *Server) GracefullyStop(ctx context.Context) error {
 	log.Infof("[HTTP] server is stopping")
 	return s.Shutdown(ctx)
 }
diff --git a/transport/transport.go b/transport/transport.go
index e315cc7..a0e5667 100644
--- a/transport/transport.go
+++ b/transport/transport.go
@@ -9,7 +9,6 @@ import (
 type Server interface {
 	Start(context.Context) error
 	Stop(context.Context) error
-	GracefullyStop(ctx context.Context) error
 }
 
 // Endpointer is registry endpoint.