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

     1  package zhttp
     2  
     3  import (
     4  	"bufio"
     5  	"crypto/tls"
     6  	"errors"
     7  	"io"
     8  	"net"
     9  	"net/http"
    10  	"net/rpc"
    11  	"net/rpc/jsonrpc"
    12  	"strconv"
    13  	"strings"
    14  	"time"
    15  
    16  	"github.com/sohaha/zlsgo/zstring"
    17  	"github.com/sohaha/zlsgo/zvalid"
    18  )
    19  
    20  type JSONRPC struct {
    21  	client  *rpc.Client
    22  	path    string
    23  	address string
    24  	options JSONRPCOptions
    25  }
    26  
    27  func (j *JSONRPC) Call(serviceMethod string, args interface{}, reply interface{}) error {
    28  	r := <-j.Go(serviceMethod, args, reply, make(chan *rpc.Call, 1)).Done
    29  
    30  	return r.Error
    31  }
    32  
    33  func (j *JSONRPC) Go(serviceMethod string, args interface{}, reply interface{}, done chan *rpc.Call) *rpc.Call {
    34  	r := j.client.Go(serviceMethod, args, reply, done)
    35  	if r.Error != nil && r.Error == rpc.ErrShutdown && j.options.Retry {
    36  		if j.options.RetryDelay > 0 {
    37  			time.Sleep(j.options.RetryDelay)
    38  		}
    39  		err := j.connect()
    40  		if err == nil {
    41  			done = make(chan *rpc.Call, 1)
    42  			return j.Go(serviceMethod, args, reply, done)
    43  		}
    44  	}
    45  	return r
    46  }
    47  
    48  func (j *JSONRPC) Close() error {
    49  	err := j.client.Close()
    50  	return err
    51  }
    52  
    53  func (j *JSONRPC) connect() error {
    54  	var conn net.Conn
    55  	var err error
    56  
    57  	d := net.Dialer{}
    58  	if j.options.Timeout > 0 {
    59  		d.Timeout = j.options.Timeout
    60  	}
    61  
    62  	if j.options.TlsConfig == nil {
    63  		conn, err = d.Dial("tcp", j.address)
    64  	} else {
    65  		config := j.options.TlsConfig
    66  		if config.RootCAs == nil {
    67  			config.InsecureSkipVerify = true
    68  		}
    69  		conn, err = tls.DialWithDialer(&d, "tcp", j.address, config)
    70  	}
    71  
    72  	if err != nil {
    73  		return err
    74  	}
    75  
    76  	message := zstring.Buffer(2)
    77  	_, _ = message.WriteString("CONNECT " + j.path + " HTTP/1.0\n")
    78  
    79  	for k := range j.options.Header {
    80  		_, _ = message.WriteString(k)
    81  		_, _ = message.WriteString(": ")
    82  		_, _ = message.WriteString(j.options.Header.Get(k))
    83  		_, _ = message.WriteString("\n")
    84  	}
    85  
    86  	_, _ = message.WriteString("\n\n")
    87  	_, _ = io.WriteString(conn, message.String())
    88  	response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: "CONNECT"})
    89  	if err != nil {
    90  		return err
    91  	}
    92  
    93  	if response.StatusCode != http.StatusOK && response.ContentLength != -1 {
    94  		return errors.New("Prohibit connection, a status code: " + strconv.Itoa(response.StatusCode))
    95  	}
    96  	j.client = rpc.NewClientWithCodec(jsonrpc.NewClientCodec(conn))
    97  
    98  	return nil
    99  }
   100  
   101  type JSONRPCOptions struct {
   102  	TlsConfig  *tls.Config
   103  	Header     http.Header
   104  	Timeout    time.Duration
   105  	RetryDelay time.Duration
   106  	Retry      bool
   107  }
   108  
   109  func NewJSONRPC(address string, path string, opts ...func(o *JSONRPCOptions)) (client *JSONRPC, err error) {
   110  	o := JSONRPCOptions{
   111  		Retry:      true,
   112  		RetryDelay: time.Second * 1,
   113  		Header:     http.Header{},
   114  	}
   115  	if len(opts) > 0 {
   116  		opts[0](&o)
   117  	}
   118  
   119  	if !strings.ContainsRune(address, ':') {
   120  		address = ":" + address
   121  	}
   122  
   123  	if zvalid.Text(address).IsURL().Ok() {
   124  		s := strings.Split(address, "://")
   125  		if len(s) > 1 {
   126  			address = s[1]
   127  			if s[0] == "https" {
   128  				o.TlsConfig = &tls.Config{}
   129  			}
   130  		}
   131  	}
   132  
   133  	client = &JSONRPC{options: o, address: address, path: path}
   134  	err = client.connect()
   135  	return
   136  }