refactor the code related to xtcp (#3449)

pkg/transport/message.go

@@ -0,0 +1,119 @@
+// Copyright 2023 The frp Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package transport
+
+import (
+	"context"
+	"reflect"
+	"sync"
+
+	"github.com/fatedier/golib/errors"
+
+	"github.com/fatedier/frp/pkg/msg"
+)
+
+type MessageTransporter interface {
+	Send(msg.Message) error
+	// Recv(ctx context.Context, laneKey string, msgType string) (Message, error)
+	// Do will first send msg, then recv msg with the same laneKey and specified msgType.
+	Do(ctx context.Context, req msg.Message, laneKey, recvMsgType string) (msg.Message, error)
+	Dispatch(m msg.Message, laneKey string) bool
+	DispatchWithType(m msg.Message, msgType, laneKey string) bool
+}
+
+func NewMessageTransporter(sendCh chan msg.Message) MessageTransporter {
+	return &transporterImpl{
+		sendCh:   sendCh,
+		registry: make(map[string]map[string]chan msg.Message),
+	}
+}
+
+type transporterImpl struct {
+	sendCh chan msg.Message
+
+	// First key is message type and second key is lane key.
+	// Dispatch will dispatch message to releated channel by its message type
+	// and lane key.
+	registry map[string]map[string]chan msg.Message
+	mu       sync.RWMutex
+}
+
+func (impl *transporterImpl) Send(m msg.Message) error {
+	return errors.PanicToError(func() {
+		impl.sendCh <- m
+	})
+}
+
+func (impl *transporterImpl) Do(ctx context.Context, req msg.Message, laneKey, recvMsgType string) (msg.Message, error) {
+	ch := make(chan msg.Message, 1)
+	defer close(ch)
+	unregisterFn := impl.registerMsgChan(ch, laneKey, recvMsgType)
+	defer unregisterFn()
+
+	if err := impl.Send(req); err != nil {
+		return nil, err
+	}
+
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case resp := <-ch:
+		return resp, nil
+	}
+}
+
+func (impl *transporterImpl) DispatchWithType(m msg.Message, msgType, laneKey string) bool {
+	var ch chan msg.Message
+	impl.mu.RLock()
+	byLaneKey, ok := impl.registry[msgType]
+	if ok {
+		ch = byLaneKey[laneKey]
+	}
+	impl.mu.RUnlock()
+
+	if ch == nil {
+		return false
+	}
+
+	if err := errors.PanicToError(func() {
+		ch <- m
+	}); err != nil {
+		return false
+	}
+	return true
+}
+
+func (impl *transporterImpl) Dispatch(m msg.Message, laneKey string) bool {
+	msgType := reflect.TypeOf(m).Elem().Name()
+	return impl.DispatchWithType(m, msgType, laneKey)
+}
+
+func (impl *transporterImpl) registerMsgChan(recvCh chan msg.Message, laneKey string, msgType string) (unregister func()) {
+	impl.mu.Lock()
+	byLaneKey, ok := impl.registry[msgType]
+	if !ok {
+		byLaneKey = make(map[string]chan msg.Message)
+		impl.registry[msgType] = byLaneKey
+	}
+	byLaneKey[laneKey] = recvCh
+	impl.mu.Unlock()
+
+	unregister = func() {
+		impl.mu.Lock()
+		delete(byLaneKey, laneKey)
+		impl.mu.Unlock()
+	}
+	return
+}

pkg/transport/tls.go

@@ -1,3 +1,17 @@
+// Copyright 2023 The frp Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 package transport
 
 import (