AutomatedTesting/usecases_server/utils/cyllib/socket/socket.go

package sorket

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"

	"net"
	"net/http"
	"time"
	"automatedtesting/usecases_server/utils/cyllib/log"

	"github.com/gorilla/websocket"
)

// socket
var upgrader = websocket.Upgrader{
	// 解决跨域问题
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
} // use default options

type WebSocket struct {
	ws  *websocket.Conn
	mux sync.Mutex
}

func NewWebSocket() *WebSocket {
	sc := WebSocket{}
	return &sc
}
func (s *WebSocket) Connect(w http.ResponseWriter, r *http.Request) bool {
	defer func() {
		if err := recover(); err != nil {
			if s.ws != nil {
				s.ws.Close()
			}
		}
	}()
	wsw, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Println(err)
		return false
	}
	s.ws = wsw
	return true
}
func (s *WebSocket) Close() bool {
	return s.ws.Close() == nil
}

func (s *WebSocket) ReadBytes() ([]byte, bool) {
	_, data, err := s.ws.ReadMessage()
	if err != nil {
		return nil, false
	}
	return data, true
}

func (s *WebSocket) SendJson(data interface{}) bool {
	byt, err := json.Marshal(data)
	if err != nil {
		return false
	}
	return s.SendBytes(byt)
}
func (s *WebSocket) SendBytes(data []byte) bool {
	s.mux.Lock()
	err := s.ws.WriteMessage(1, data)
	s.mux.Unlock()
	return err == nil
}

// ---------------------------------------------------------------
// ---------------------------------------------------------------

const (
	TYPE_TCP = "tcp"
	TYPE_UDP = "udp"
)

type SocKetOsi struct {
	typ        string
	socket     net.Listener
	conn       net.Conn
	contimeout int
}

func NewSocket(typ string, connect_timeout int) *SocKetOsi {
	sc := SocKetOsi{}
	sc.typ = typ
	sc.contimeout = connect_timeout
	return &sc
}

func (s *SocKetOsi) Connect(ip string) (bool, error) {
	var err error
	s.socket, err = net.Listen(s.typ, ":9101")
	if err != nil {
		return false, err
	}
	defer s.socket.Close()
	ctx, _ := context.WithTimeout(context.Background(), time.Duration(s.contimeout)*time.Millisecond)
	flag := false
	go func(flag *bool) {
		s.conn, err = s.socket.Accept()
		if err != nil {
			*flag = false
			return
		}
		*flag = true
	}(&flag)
Loop2:
	for {
		select {
		case <-ctx.Done():
			s.socket.Close()
			log.Println("超时时间到了")
			break Loop2
		default:
			if flag {
				break Loop2
			}
		}
	}
	df := func() {
		if s.conn != nil {
			s.conn.Close()
		}
	}
	defer df()
	return flag, nil
}

func (s *SocKetOsi) Close() {
	if s.conn != nil {
		s.conn.Close()
	}
	if s.socket != nil {
		s.socket.Close()
	}
}
func (s *SocKetOsi) WriteBytes(data []byte) bool {
	_, err := s.conn.Write(data)
	log.WriteErr(err)
	return err == nil
}
func (s *SocKetOsi) ReadBytes() ([]byte, bool) {
	var buf = make([]byte, 1024)
	_, err := s.conn.Read(buf)
	if err != nil {
		return nil, false
	}
	return buf, true
}
func (s *SocKetOsi) WriteData(data interface{}) bool {
	byt, err := json.Marshal(data)
	if err != nil {
		return false
	}
	return s.WriteBytes(byt)
}