Included .env file and #30 Modbus emulation

ryabenbc · ryabenbc · commit ed1b2b17cf0e · 2021-05-11T16:04:14.000+02:00
diff --git a/.gitignore b/.gitignore
@@ -24,11 +24,6 @@ bin/
 ## VSCode workspace file
 workspace*.code-*
 
-# environment files (not example)
-.env
-!.env.example*
-!*.example
-
 # Vendor folder
 /vendor/
 
diff --git a/build/docker/Dockerfile b/build/docker/Dockerfile
@@ -16,7 +16,7 @@ COPY go.sum .
 # download all the dependencies
 RUN go mod download
 
-# install netcat
+# install netcat so we can ping the db until it is ready
 RUN apt-get update && apt-get install -y netcat
 
 # Copy everything into the image
@@ -29,8 +29,8 @@ COPY . .
 # Disclaimer: if you comment this line, be 100% sure that the binary can be run on linux
 RUN make riotpot-builder
 
+# give permissions to the entrypoint to run the file
 RUN chmod +x build/docker/entrypoint.sh
 
 # Run RIoTPot
 ENTRYPOINT [ "./build/docker/entrypoint.sh" ]
-#CMD ["./riotpot"]
diff --git a/build/env/.env b/build/env/.env
@@ -0,0 +1,6 @@
+AUTOD=false
+DB_HOST=postgres
+DB_USER=user
+DB_PASS=password
+DB_NAME=db
+DB_PORT=5432
diff --git a/pkg/plugin/modbusd/modbusd.go b/pkg/plugin/modbusd/modbusd.go
@@ -2,8 +2,10 @@ package main
 
 import (
 	"fmt"
+	"io"
 	"net"
 
+	"github.com/riotpot/pkg/models"
 	"github.com/riotpot/pkg/services"
 	"github.com/riotpot/tools/errors"
 	"github.com/xiegeo/modbusone"
@@ -32,16 +34,23 @@ func Modbusd() services.Service {
 		Protocol: "tcp",
 	}
 
+	handler := handler()
+
 	return &Modbus{
 		mx,
+		handler,
 	}
 }
 
 type Modbus struct {
 	services.MixinService
+	handler modbusone.ProtocolHandler
 }
 
 func (m *Modbus) Run() (err error) {
+	// before running, migrate the model that we want to store
+	m.Migrate(&models.Connection{})
+
 	// convert the port number to a string that we can use it in the server
 	var port = fmt.Sprintf(":%d", m.Port)
 
@@ -52,35 +61,134 @@ func (m *Modbus) Run() (err error) {
 	// create the channel for stopping the service
 	m.StopCh = make(chan int, 1)
 
-	// serve the app
-	server := modbusone.NewTCPServer(listener)
-	m.serve(server)
+	// build a channel stack to receive connections to the service
+	conn := make(chan net.Conn)
+	m.serve(conn, listener)
 
 	// update the status of the service
 	m.Running <- true
 
 	// block here until we receive an stopping signal in the channel
 	<-m.StopCh
 
-	// close the server
-	defer server.Close()
-
 	// Close the channel for stopping the service
 	fmt.Print("[x] Service stopped...\n")
 	close(m.StopCh)
 
 	return
 }
 
-func (m *Modbus) serve(server modbusone.ServerCloser) {
+// Open the service and listen for connections
+// inspired on https://gist.github.com/paulsmith/775764#file-echo-go
+func (m *Modbus) serve(ch chan net.Conn, listener net.Listener) {
+	// open an infinite loop to receive connections
 	fmt.Printf("[%s] Started listenning for connections in port %d\n", Name, m.Port)
+	for {
+		// Accept the client connection
+		client, err := listener.Accept()
+		if err != nil {
+			return
+		}
+		defer client.Close()
+
+		// push the client connection to the channel
+		ch <- client
+	}
+}
+
+// Handle the pool of connections to the service
+func (m *Modbus) handlePool(ch chan net.Conn) {
+	// open an infinite loop to handle the connections
+	for {
+		// while the `stop` channel remains empty, continue handling
+		// new connections.
+		select {
+		case <-m.StopCh:
+			// stop the pool
+			fmt.Printf("[x] Stopping %s service...\n", m.Name)
+			// update the status of the service
+			m.Running <- false
+			return
+		case conn := <-ch:
+			// use one goroutine per connection.
+			go m.handleSession(conn)
+		}
+	}
+}
+
+// Handle a connection made to the service
+// We are rewriting this function so we can interact with the connection loop
+// and rewrite some of the underlying function. Unfortunately, the loop
+// is deep and contains a goroutine in it as main handler of the session connections.
+func (m *Modbus) handleSession(conn net.Conn) {
+	for {
+		defer conn.Close()
+
+		wec := func(conn net.Conn, bs []byte, req modbusone.PDU, err error) {
+			writeTCP(conn, bs, modbusone.ExceptionReplyPacket(req, modbusone.ToExceptionCode(err)))
+		}
+
+		var rb []byte
+		if modbusone.OverSizeSupport {
+			rb = make(
+				[]byte,
+				modbusone.MBAPHeaderLength+modbusone.OverSizeMaxRTU+modbusone.TCPHeaderLength,
+			)
+		} else {
+			rb = make([]byte, modbusone.MBAPHeaderLength+modbusone.MaxPDUSize)
+		}
+
+		for {
+			n, err := readTCP(conn, rb)
+			if err != nil {
+				return
+			}
+
+			// load the payload from the packet i.e. everything after the header
+			p := modbusone.PDU(rb[modbusone.MBAPHeaderLength:n])
+
+			// save the request connection with the request payload in it.
+			m.save(conn, p)
+
+			// validate the request, checks for errors on the code and the
+			// length of the payload
+			err = p.ValidateRequest()
+			if err != nil {
+				m.save(conn, p)
+				return
+			}
+
+			fc := p.GetFunctionCode()
+
+			// initialize the data. It will be filled with the information
+			// in the payload.
+			var data []byte
+
+			// Only two things can happen,
+			// either read from the server or write to it.
+			if fc.IsReadToServer() {
+				data, err = m.handler.OnRead(p)
+				if err != nil {
+					wec(conn, rb, p, err)
+					continue
+				}
+				writeTCP(conn, rb, p.MakeReadReply(data))
+			} else if fc.IsWriteToServer() {
+				data, err = p.GetRequestValues()
+				if err != nil {
+					wec(conn, rb, p, err)
+					continue
+				}
+				err = m.handler.OnWrite(p, data)
+				if err != nil {
+					wec(conn, rb, p, err)
+					continue
+				}
+				writeTCP(conn, rb, p.MakeWriteReply())
+			}
 
-	// since we don't have to handle the connection as the server already
-	// has a handler, we just run the server in another goroutine
-	go func() {
-		err := server.Serve(m.handler("server"))
-		errors.Raise(err)
-	}()
+		}
+	}
 }
 
 // Simple handler for the Modbus functions.
@@ -97,7 +205,7 @@ func (m *Modbus) serve(server modbusone.ServerCloser) {
 // 	16: multiple 6
 // Generarly, we will either read a quantity of an unsigned int 16, or write
 // booleans (0 or 1) plus the address.
-func (m *Modbus) handler(name string) modbusone.ProtocolHandler {
+func handler() modbusone.ProtocolHandler {
 	return &modbusone.SimpleHandler{
 
 		// Discrete Inputs
@@ -151,3 +259,47 @@ func (m *Modbus) handler(name string) modbusone.ProtocolHandler {
 		},
 	}
 }
+
+// Copy pasted from https://github.com/xiegeo/modbusone/blob/797d647e237d97ab9d2bdad49bf42591ea7076f2/tcp_server.go#L19
+// both functions are just packet handlers that read the content and headers.
+// It is unnecessary to reimplement something already done, however, this function
+// are now exported but still necessary for our connection handler.
+func readTCP(r io.Reader, bs []byte) (n int, err error) {
+	n, err = io.ReadFull(r, bs[:modbusone.TCPHeaderLength])
+	if err != nil {
+		return n, err
+	}
+	if bs[2] != 0 || bs[3] != 0 {
+		return n, fmt.Errorf("MBAP protocol of %X %X is unknown", bs[2], bs[3])
+	}
+	l := int(bs[4])*256 + int(bs[5])
+	if l <= 2 {
+		return n, fmt.Errorf("MBAP data length of %v is too short, bs:%x", l, bs[:n])
+	}
+	if len(bs) < l+modbusone.TCPHeaderLength {
+		return n, fmt.Errorf("MBAP data length of %v is too long", l)
+	}
+	n, err = io.ReadFull(r, bs[modbusone.TCPHeaderLength:l+modbusone.TCPHeaderLength])
+	return n + modbusone.TCPHeaderLength, err
+}
+
+//writeTCP writes a PDU packet on TCP reusing the headers and buffer space in bs
+func writeTCP(w io.Writer, bs []byte, pdu modbusone.PDU) (int, error) {
+	l := len(pdu) + 1 //pdu + byte of slaveID
+	bs[4] = byte(l / 256)
+	bs[5] = byte(l)
+	copy(bs[modbusone.MBAPHeaderLength:], pdu)
+	return w.Write(bs[:len(pdu)+modbusone.MBAPHeaderLength])
+}
+
+func (m *Modbus) save(conn net.Conn, payload []byte) {
+	connection := models.NewConnection()
+	connection.LocalAddress = conn.LocalAddr().String()
+	connection.RemoteAddress = conn.RemoteAddr().String()
+	connection.Protocol = "TCP"
+	connection.Service = Name
+	connection.Incoming = true
+	connection.Payload = string(payload)
+
+	m.Store(connection)
+}
