github.com/Kong/go-pdk@v0.11.0/bridge/bridge.go (about)

     1  /*
     2  Used internally for the RPC protocol.
     3  */
     4  package bridge
     5  
     6  import (
     7  	"encoding/binary"
     8  	"errors"
     9  	"io"
    10  	"log"
    11  	"net"
    12  
    13  	"github.com/Kong/go-pdk/server/kong_plugin_protocol"
    14  	"google.golang.org/protobuf/proto"
    15  	"google.golang.org/protobuf/types/known/structpb"
    16  )
    17  
    18  type PdkBridge struct {
    19  	conn net.Conn
    20  }
    21  
    22  type StepData struct {
    23  	Method string
    24  	Args   []interface{}
    25  }
    26  
    27  func New(conn net.Conn) PdkBridge {
    28  	return PdkBridge{
    29  		conn: conn,
    30  	}
    31  }
    32  
    33  func readPbFrame(conn net.Conn) (data []byte, err error) {
    34  	var len uint32
    35  	err = binary.Read(conn, binary.LittleEndian, &len)
    36  	if err != nil {
    37  		return
    38  	}
    39  
    40  	data = make([]byte, len)
    41  
    42  	_, err = io.ReadFull(conn, data)
    43  	if err != nil {
    44  		return nil, err
    45  	}
    46  
    47  	return
    48  }
    49  
    50  func writePbFrame(conn net.Conn, data []byte) (err error) {
    51  	var len uint32 = uint32(len(data))
    52  	err = binary.Write(conn, binary.LittleEndian, len)
    53  	if err != nil {
    54  		return
    55  	}
    56  
    57  	if len > 0 {
    58  		_, err = conn.Write(data)
    59  	}
    60  
    61  	return
    62  }
    63  
    64  func WrapString(s string) *kong_plugin_protocol.String {
    65  	return &kong_plugin_protocol.String{V: s}
    66  }
    67  
    68  func WrapByteString(s []byte) *kong_plugin_protocol.ByteString {
    69  	return &kong_plugin_protocol.ByteString{V: s}
    70  }
    71  
    72  func WrapHeaders(h map[string][]string) (*structpb.Struct, error) {
    73  	h2 := make(map[string]interface{}, len(h))
    74  	for k, v := range h {
    75  		l := make([]interface{}, len(v))
    76  		for i, v2 := range v {
    77  			l[i] = v2
    78  		}
    79  		h2[k] = l
    80  	}
    81  
    82  	st, err := structpb.NewStruct(h2)
    83  	if err != nil {
    84  		return nil, err
    85  	}
    86  
    87  	return st, nil
    88  }
    89  
    90  func UnwrapHeaders(st *structpb.Struct) map[string][]string {
    91  	m := st.AsMap()
    92  	m2 := make(map[string][]string)
    93  	for k, v := range m {
    94  		switch v2 := v.(type) {
    95  		case string:
    96  			m2[k] = []string{v2}
    97  		case []string:
    98  			m2[k] = v2
    99  		case []interface{}:
   100  			m2[k] = make([]string, len(v2))
   101  			for i, v3 := range v2 {
   102  				if s, ok := v3.(string); ok {
   103  					m2[k][i] = s
   104  				}
   105  			}
   106  		default:
   107  			log.Printf("unexpected type %T on header %s:%v", v2, k, v2)
   108  		}
   109  	}
   110  
   111  	return m2
   112  }
   113  
   114  func (b PdkBridge) Ask(method string, args proto.Message, out proto.Message) error {
   115  	err := writePbFrame(b.conn, []byte(method))
   116  	if err != nil {
   117  		return err
   118  	}
   119  
   120  	var args_d []byte
   121  
   122  	if args != nil {
   123  		args_d, err = proto.Marshal(args)
   124  		if err != nil {
   125  			return err
   126  		}
   127  	}
   128  
   129  	err = writePbFrame(b.conn, args_d)
   130  	if err != nil {
   131  		return err
   132  	}
   133  
   134  	out_d, err := readPbFrame(b.conn)
   135  	if err != nil {
   136  		return err
   137  	}
   138  
   139  	if out != nil {
   140  		err = proto.Unmarshal(out_d, out)
   141  	}
   142  
   143  	return err
   144  }
   145  
   146  func (b PdkBridge) AskString(method string, args proto.Message) (string, error) {
   147  	out := new(kong_plugin_protocol.String)
   148  	err := b.Ask(method, args, out)
   149  	return out.V, err
   150  }
   151  
   152  func (b PdkBridge) AskInt(method string, args proto.Message) (int, error) {
   153  	out := new(kong_plugin_protocol.Int)
   154  	err := b.Ask(method, args, out)
   155  	return int(out.V), err
   156  }
   157  
   158  func (b PdkBridge) AskNumber(method string, args proto.Message) (float64, error) {
   159  	out := new(kong_plugin_protocol.Number)
   160  	err := b.Ask(method, args, out)
   161  	return out.V, err
   162  }
   163  
   164  func (b PdkBridge) AskValue(method string, args proto.Message) (interface{}, error) {
   165  	out := new(structpb.Value)
   166  	err := b.Ask(method, args, out)
   167  	if err != nil {
   168  		return nil, err
   169  	}
   170  
   171  	return out.AsInterface(), nil
   172  }
   173  
   174  func (b PdkBridge) Close() error {
   175  	return b.conn.Close()
   176  }
   177  
   178  func ReturnTypeError(expected string) error {
   179  	return errors.New("expected type: " + expected)
   180  }