github.com/sohaha/zlsgo@v1.7.13-0.20240501141223-10dd1a906f76/znet/rpc.go (about)

     1  package znet
     2  
     3  import (
     4  	"io"
     5  	"net/http"
     6  	"net/rpc"
     7  	"net/rpc/jsonrpc"
     8  
     9  	"github.com/sohaha/zlsgo/zreflect"
    10  )
    11  
    12  type JSONRPCOption struct {
    13  	DisabledHTTP bool
    14  	Debug        bool
    15  }
    16  
    17  func JSONRPC(rcvr map[string]interface{}, opts ...func(o *JSONRPCOption)) func(c *Context) {
    18  	o := JSONRPCOption{}
    19  	if len(opts) > 0 {
    20  		opts[0](&o)
    21  	}
    22  
    23  	s := rpc.NewServer()
    24  	methods := make(map[string][]string, 0)
    25  	for name, v := range rcvr {
    26  		err := s.RegisterName(name, v)
    27  		if err == nil && o.Debug {
    28  			typ := zreflect.TypeOf(v)
    29  			for m := 0; m < typ.NumMethod(); m++ {
    30  				method := typ.Method(m)
    31  				mtype := method.Type
    32  				mname := method.Name
    33  				l := mtype.NumIn()
    34  				replyType, argType := "-", "-"
    35  				if l > 2 {
    36  					replyType = mtype.In(2).String()
    37  				}
    38  				if l > 1 {
    39  					argType = mtype.In(1).String()
    40  				}
    41  				methods[name+"."+mname] = []string{argType, replyType}
    42  			}
    43  		}
    44  	}
    45  
    46  	return func(c *Context) {
    47  		req := c.Request
    48  		method := req.Method
    49  		if o.Debug && method == "GET" {
    50  			c.JSON(200, methods)
    51  			return
    52  		}
    53  
    54  		if c.stopHandle.Load() {
    55  			return
    56  		}
    57  
    58  		var codec rpc.ServerCodec
    59  		if method == "CONNECT" || (method == "POST" && !o.DisabledHTTP) {
    60  			c.stopHandle.Store(true)
    61  			c.write()
    62  
    63  			if method == "CONNECT" {
    64  				conn, _, _ := c.Writer.(http.Hijacker).Hijack()
    65  				codec = jsonrpc.NewServerCodec(conn)
    66  				_, _ = io.WriteString(conn, "HTTP/1.0 200 Connected to JSON RPC\n\n")
    67  				s.ServeCodec(codec)
    68  				return
    69  			}
    70  
    71  			c.Writer.Header().Set("Content-Type", ContentTypeJSON)
    72  			var conn io.ReadWriteCloser = struct {
    73  				io.Writer
    74  				io.ReadCloser
    75  			}{
    76  				ReadCloser: c.Request.Body,
    77  				Writer:     c.Writer,
    78  			}
    79  			_ = s.ServeRequest(jsonrpc.NewServerCodec(conn))
    80  			return
    81  		}
    82  
    83  		c.SetContentType(ContentTypePlain)
    84  		c.String(http.StatusMethodNotAllowed, "405 must CONNECT\n")
    85  	}
    86  }