github.com/cloudwego/kitex@v0.9.0/pkg/remote/codec/bytebuf_util.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package codec
    18  
    19  import (
    20  	"encoding/binary"
    21  	"errors"
    22  	"io"
    23  
    24  	"github.com/cloudwego/kitex/pkg/remote"
    25  )
    26  
    27  // ReadUint32 ...
    28  func ReadUint32(in remote.ByteBuffer) (uint32, error) {
    29  	var buf []byte
    30  	var err error
    31  	if buf, err = in.Next(Size32); err != nil {
    32  		return 0, err
    33  	}
    34  	return binary.BigEndian.Uint32(buf), nil
    35  }
    36  
    37  // PeekUint32 ...
    38  func PeekUint32(in remote.ByteBuffer) (uint32, error) {
    39  	var buf []byte
    40  	var err error
    41  	if buf, err = in.Peek(Size32); err != nil {
    42  		return 0, err
    43  	}
    44  	return binary.BigEndian.Uint32(buf), nil
    45  }
    46  
    47  // ReadUint16 ...
    48  func ReadUint16(in remote.ByteBuffer) (uint16, error) {
    49  	buf, err := in.Next(Size16)
    50  	if err != nil {
    51  		return 0, err
    52  	}
    53  	n := binary.BigEndian.Uint16(buf)
    54  	return n, nil
    55  }
    56  
    57  // WriteUint32 ...
    58  func WriteUint32(val uint32, out remote.ByteBuffer) error {
    59  	var buf []byte
    60  	var err error
    61  	if buf, err = out.Malloc(Size32); err != nil {
    62  		return err
    63  	}
    64  	binary.BigEndian.PutUint32(buf, val)
    65  	return nil
    66  }
    67  
    68  // WriteUint16 ...
    69  func WriteUint16(val uint16, out remote.ByteBuffer) error {
    70  	var buf []byte
    71  	var err error
    72  	if buf, err = out.Malloc(Size16); err != nil {
    73  		return err
    74  	}
    75  	binary.BigEndian.PutUint16(buf, val)
    76  	return nil
    77  }
    78  
    79  // WriteByte ...
    80  func WriteByte(val byte, out remote.ByteBuffer) error {
    81  	var buf []byte
    82  	var err error
    83  	if buf, err = out.Malloc(1); err != nil {
    84  		return err
    85  	}
    86  	buf[0] = val
    87  	return nil
    88  }
    89  
    90  // Bytes2Uint32NoCheck ...
    91  func Bytes2Uint32NoCheck(bytes []byte) uint32 {
    92  	return binary.BigEndian.Uint32(bytes)
    93  }
    94  
    95  // Bytes2Uint32 ...
    96  func Bytes2Uint32(bytes []byte) (uint32, error) {
    97  	if len(bytes) < 4 {
    98  		return 0, io.EOF
    99  	}
   100  	return binary.BigEndian.Uint32(bytes), nil
   101  }
   102  
   103  // Bytes2Uint16NoCheck ...
   104  func Bytes2Uint16NoCheck(bytes []byte) uint16 {
   105  	return binary.BigEndian.Uint16(bytes)
   106  }
   107  
   108  // Bytes2Uint16 ...
   109  func Bytes2Uint16(bytes []byte, off int) (uint16, error) {
   110  	if len(bytes)-off < 2 {
   111  		return 0, io.EOF
   112  	}
   113  	return binary.BigEndian.Uint16(bytes[off:]), nil
   114  }
   115  
   116  // Bytes2Uint8 ...
   117  func Bytes2Uint8(bytes []byte, off int) (uint8, error) {
   118  	if len(bytes)-off < 1 {
   119  		return 0, io.EOF
   120  	}
   121  	return bytes[off], nil
   122  }
   123  
   124  // ReadString ...
   125  func ReadString(in remote.ByteBuffer) (string, int, error) {
   126  	strLen, err := ReadUint32(in)
   127  	if err != nil {
   128  		return "", 0, err
   129  	}
   130  	if strLen == 0 {
   131  		return "", 0, errors.New("invalid string length in ReadString")
   132  	}
   133  	str, err := in.ReadString(int(strLen))
   134  	return str, int(strLen) + Size32, err
   135  }
   136  
   137  // WriteString ...
   138  func WriteString(val string, out remote.ByteBuffer) (int, error) {
   139  	strLen := len(val)
   140  	if err := WriteUint32(uint32(strLen), out); err != nil {
   141  		return 0, err
   142  	}
   143  	n, err := out.WriteString(val)
   144  	if err != nil {
   145  		return 0, err
   146  	}
   147  	return n + 4, nil
   148  }
   149  
   150  // WriteString2BLen ...
   151  func WriteString2BLen(val string, out remote.ByteBuffer) (int, error) {
   152  	strLen := len(val)
   153  	if err := WriteUint16(uint16(strLen), out); err != nil {
   154  		return 0, err
   155  	}
   156  	n, err := out.WriteString(val)
   157  	if err != nil {
   158  		return 0, err
   159  	}
   160  	return n + 2, nil
   161  }
   162  
   163  // ReadString2BLen ...
   164  func ReadString2BLen(bytes []byte, off int) (string, int, error) {
   165  	length, err := Bytes2Uint16(bytes, off)
   166  	strLen := int(length)
   167  	if err != nil {
   168  		return "", 0, err
   169  	}
   170  	off += 2
   171  	if len(bytes)-off < strLen {
   172  		return "", 0, io.EOF
   173  	}
   174  
   175  	buf := bytes[off : off+strLen]
   176  	return string(buf), int(length) + 2, nil
   177  }