github.com/jingruilea/kubeedge@v1.2.0-beta.0.0.20200410162146-4bb8902b3879/cloud/pkg/cloudhub/servers/udsserver/uds.go (about)

     1  package udsserver
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"os"
     7  	"strings"
     8  
     9  	"k8s.io/klog"
    10  )
    11  
    12  const (
    13  	// DefaultBufferSize represents default buffer size
    14  	DefaultBufferSize = 10480
    15  )
    16  
    17  // UnixDomainSocket struct
    18  type UnixDomainSocket struct {
    19  	filename   string
    20  	buffersize int
    21  	handler    func(string) string
    22  }
    23  
    24  // NewUnixDomainSocket create new socket
    25  func NewUnixDomainSocket(filename string, buffersize ...int) *UnixDomainSocket {
    26  	size := DefaultBufferSize
    27  	if buffersize != nil {
    28  		size = buffersize[0]
    29  	}
    30  	us := UnixDomainSocket{filename: filename, buffersize: size}
    31  	return &us
    32  }
    33  
    34  // parseEndpoint parses endpoint
    35  func parseEndpoint(ep string) (string, string, error) {
    36  	if strings.HasPrefix(strings.ToLower(ep), "unix://") || strings.HasPrefix(strings.ToLower(ep), "tcp://") {
    37  		s := strings.SplitN(ep, "://", 2)
    38  		if s[1] != "" {
    39  			return s[0], s[1], nil
    40  		}
    41  	}
    42  	return "", "", fmt.Errorf("invalid endpoint: %v", ep)
    43  }
    44  
    45  // SetContextHandler set handler for server
    46  func (us *UnixDomainSocket) SetContextHandler(f func(string) string) {
    47  	us.handler = f
    48  }
    49  
    50  // StartServer start for server
    51  func (us *UnixDomainSocket) StartServer() error {
    52  	proto, addr, err := parseEndpoint(us.filename)
    53  	if err != nil {
    54  		klog.Errorf("failed to parseEndpoint: %v", err)
    55  		return err
    56  	}
    57  	if proto == "unix" {
    58  		addr = "/" + addr
    59  		if err := os.Remove(addr); err != nil && !os.IsNotExist(err) { //nolint: vetshadow
    60  			klog.Errorf("failed to remove addr: %v", err)
    61  			return err
    62  		}
    63  	}
    64  
    65  	// Listen
    66  	listener, err := net.Listen(proto, addr)
    67  	if err != nil {
    68  		klog.Errorf("failed to listen addr: %v", err)
    69  		return err
    70  	}
    71  	defer listener.Close()
    72  	klog.Infof("listening on: %v", listener.Addr())
    73  
    74  	for {
    75  		c, err := listener.Accept()
    76  		if err != nil {
    77  			klog.Errorf("accept to error: %v", err)
    78  			continue
    79  		}
    80  		go us.handleServerConn(c)
    81  	}
    82  }
    83  
    84  // handleServerConn handler for server
    85  func (us *UnixDomainSocket) handleServerConn(c net.Conn) {
    86  	defer c.Close()
    87  	buf := make([]byte, us.buffersize)
    88  	nr, err := c.Read(buf)
    89  	if err != nil {
    90  		klog.Errorf("failed to read buffer: %v", err)
    91  		return
    92  	}
    93  	result := us.handleServerContext(string(buf[0:nr]))
    94  	_, err = c.Write([]byte(result))
    95  	if err != nil {
    96  		klog.Errorf("failed to write buffer: %v", err)
    97  	}
    98  }
    99  
   100  // HandleServerContext handler for server
   101  func (us *UnixDomainSocket) handleServerContext(context string) string {
   102  	if us.handler != nil {
   103  		return us.handler(context)
   104  	}
   105  	return ""
   106  }