feat: support config MaxResponseSize (#46)

Signed-off-by: ZhangJian He <shoothzj@gmail.com>
protocol-laboratory · Oct 28, 2024 · 3af6e5d · 3af6e5d
1 parent cdb2753
commit 3af6e5d
Show file tree

Hide file tree

Showing 8 changed files with 40 additions and 13 deletions.
diff --git a/opcua/enc/encoder.go b/opcua/enc/encoder.go
@@ -24,10 +24,13 @@ type FastEncoder interface {
 
 type DefaultEncoder struct {
 	sequenceNumberGenerator func() uint32
+	maxEncodedSize          int
 }
 
-func NewDefaultEncoder() *DefaultEncoder {
-	return &DefaultEncoder{}
+func NewDefaultEncoder(maxEncodedSize int) *DefaultEncoder {
+	return &DefaultEncoder{
+		maxEncodedSize: maxEncodedSize,
+	}
 }
 
 func (e *DefaultEncoder) Encode(v *uamsg.Message, chunkSize int) ([][]byte, error) {
@@ -67,6 +70,7 @@ func (e *DefaultEncoder) Encode(v *uamsg.Message, chunkSize int) ([][]byte, erro
 
 	leftBodySize := len(dataBytes)
 	headerLength := messageHeaderLength + securityHeaderLength + sequenceHeaderLength
+	totalSize := 0
 
 	for leftBodySize > 0 {
 		tempBuff := bytes.NewBuffer(nil)
@@ -120,6 +124,11 @@ func (e *DefaultEncoder) Encode(v *uamsg.Message, chunkSize int) ([][]byte, erro
 				return nil, err
 			}
 		}
+
+		totalSize += tempBuff.Len()
+		if totalSize > e.maxEncodedSize {
+			return nil, errors.New("exceed max encoded size")
+		}
 		chunks = append(chunks, tempBuff.Bytes())
 	}
 	return chunks, nil
@@ -167,6 +176,9 @@ func genericEncoder(v interface{}) ([]byte, error) {
 		if err != nil {
 			return nil, err
 		}
+		if length == 0 || length == -1 {
+			return buff.Bytes(), nil
+		}
 
 		for i := 0; i < value.Len(); i++ {
 			elemBytes, err := genericEncoder(value.Index(i).Interface())

diff --git a/opcua/enc/encoder_test.go b/opcua/enc/encoder_test.go
@@ -70,7 +70,7 @@ func showByteSlice(s []byte) {
 func getHelloMsgTestCase() *encodeTestCase {
 	return &encodeTestCase{
 		name: "encode hello msg",
-		e:    &DefaultEncoder{},
+		e:    NewDefaultEncoder(64 * 1024),
 		args: encodeArgs{
 			v: &uamsg.Message{
 				MessageHeader: &uamsg.MessageHeader{
@@ -98,7 +98,7 @@ func getHelloMsgTestCase() *encodeTestCase {
 func getAcknowledgeMsgTestCase() *encodeTestCase {
 	return &encodeTestCase{
 		name: "encode acknowledge msg",
-		e:    &DefaultEncoder{},
+		e:    NewDefaultEncoder(64 * 1024),
 		args: encodeArgs{
 			v: &uamsg.Message{
 				MessageHeader: &uamsg.MessageHeader{
@@ -124,7 +124,7 @@ func getAcknowledgeMsgTestCase() *encodeTestCase {
 func getOpenSecureChannelRequestMsgTestCase() *encodeTestCase {
 	return &encodeTestCase{
 		name: "encode open secure channel request msg",
-		e:    &DefaultEncoder{},
+		e:    NewDefaultEncoder(64 * 1024),
 		args: encodeArgs{
 			v: &uamsg.Message{
 				MessageHeader: &uamsg.MessageHeader{
@@ -183,7 +183,7 @@ func getOpenSecureChannelRequestMsgTestCase() *encodeTestCase {
 func getOpenSecureChannelResponseMsgTestCase() *encodeTestCase {
 	return &encodeTestCase{
 		name: "encode open secure channel response msg",
-		e:    &DefaultEncoder{},
+		e:    NewDefaultEncoder(64 * 1024),
 		args: encodeArgs{
 			v: &uamsg.Message{
 				MessageHeader: &uamsg.MessageHeader{
@@ -243,7 +243,7 @@ func getOpenSecureChannelResponseMsgTestCase() *encodeTestCase {
 func getCreateSessionRequestMsgTestCase() *encodeTestCase {
 	return &encodeTestCase{
 		name: "encode create session request msg",
-		e:    &DefaultEncoder{},
+		e:    NewDefaultEncoder(64 * 1024),
 		args: encodeArgs{
 			v: &uamsg.Message{
 				MessageHeader: &uamsg.MessageHeader{
@@ -314,7 +314,7 @@ func getCreateSessionRequestMsgTestCase() *encodeTestCase {
 func getActiveSessionRequestMsgTestCase() *encodeTestCase {
 	return &encodeTestCase{
 		name: "encode active session request msg",
-		e:    &DefaultEncoder{},
+		e:    NewDefaultEncoder(64 * 1024),
 		args: encodeArgs{
 			v: &uamsg.Message{
 				MessageHeader: &uamsg.MessageHeader{
@@ -386,7 +386,7 @@ func getActiveSessionRequestMsgTestCase() *encodeTestCase {
 func getActiveSessionResponseMsgTestCase() *encodeTestCase {
 	return &encodeTestCase{
 		name: "encode active session response msg",
-		e:    &DefaultEncoder{},
+		e:    NewDefaultEncoder(64 * 1024),
 		args: encodeArgs{
 			v: &uamsg.Message{
 				MessageHeader: &uamsg.MessageHeader{

diff --git a/opcua/enc/message_acknowledge_decode_test.go b/opcua/enc/message_acknowledge_decode_test.go
@@ -2,8 +2,9 @@ package enc
 
 import (
 	"bytes"
-	"github.com/stretchr/testify/require"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestDecodeAcknowledgeMessage(t *testing.T) {

diff --git a/opcua/enc/message_close_session_request_test.go b/opcua/enc/message_close_session_request_test.go
@@ -0,0 +1,8 @@
+package enc
+
+import (
+	"testing"
+)
+
+func TestDecodeCloseSessionRequestMessage(t *testing.T) {
+}
diff --git a/opcua/enc/message_hello_decode_test.go b/opcua/enc/message_hello_decode_test.go
@@ -2,8 +2,9 @@ package enc
 
 import (
 	"bytes"
-	"github.com/stretchr/testify/require"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestDecodeHelloMessage(t *testing.T) {

diff --git a/opcua/enc/message_open_secure_channel_decode_test.go b/opcua/enc/message_open_secure_channel_decode_test.go
@@ -2,8 +2,9 @@ package enc
 
 import (
 	"bytes"
-	"github.com/stretchr/testify/require"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestDecodeOpenSecureChannelMessage(t *testing.T) {

diff --git a/opcua/secure_channel.go b/opcua/secure_channel.go
@@ -53,7 +53,7 @@ func newSecureChannel(conn *Conn,
 		handler:        handler,
 		logger:         logger,
 		decoder:        enc.NewDefaultDecoder(conn, int64(svcConf.ReceiverBufferSize)),
-		encoder:        enc.NewDefaultEncoder(),
+		encoder:        enc.NewDefaultEncoder(svcConf.MaxResponseSize),
 	}
 }
 

diff --git a/opcua/server.go b/opcua/server.go
@@ -16,6 +16,7 @@ type ServerConfig struct {
 	Handler ServerHandler
 
 	ReceiverBufferSize int
+	MaxResponseSize    int
 
 	ReadRequestNodeLimit int
 
@@ -52,6 +53,9 @@ func NewServer(config *ServerConfig) (*Server, error) {
 	if config.ReceiverBufferSize < 9 {
 		return nil, fmt.Errorf("receiver buffer size must be at least 9 bytes")
 	}
+	if config.MaxResponseSize <= 0 {
+		config.MaxResponseSize = 64 * 1024
+	}
 	server := &Server{
 		config:         config,
 		channelIdGen:   &ChannelIdGen{},