diff --git a/main.go b/main.go
index cb221f0..0b33a73 100644
--- a/main.go
+++ b/main.go
@@ -25,6 +25,12 @@ type EchoRequest struct {
 	EchoByte byte
 }
 
+type MessageHandler func(m Message) error
+
+var Handlers = map[uint32]MessageHandler{
+	EchoRequestID: handleEchoRequest,
+}
+
 func (EchoRequest) ID() uint32 {
 	return EchoRequestID
 }
@@ -36,18 +42,26 @@ func unixSocketListen() error {
 	}
 
 	for conn, err := listener.Accept(); err == nil; {
-		handleConnection(conn)
+		err = handleConnection(conn)
+		if err != nil {
+			slog.Error("Error handling connection:", err)
+		}
 	}
 
 	return nil
 }
 
-func handleConnection(conn net.Conn) {
+func handleConnection(conn net.Conn) error {
 	msgChan := readSocketMessages(conn)
 
 	for msg := range msgChan {
-		handleMessage(msg)
+		err := handleMessage(msg)
+		if err != nil {
+			return err
+		}
 	}
+
+	return nil
 }
 
 func readSocketMessages(conn net.Conn) chan Message {
@@ -123,15 +137,25 @@ func readMessage(conn net.Conn) ([]byte, error) {
 	return msgContent, nil
 }
 
-func handleMessage(msg Message) {
-	switch msg.ID() {
-	case EchoRequestID:
-		handleEchoRequest(msg.(EchoRequest))
+func handleMessage(msg Message) error {
+	handler, ok := Handlers[msg.ID()]
+	if !ok {
+		return errors.New("message handler not defined for a given ID")
 	}
+
+	err := handler(msg)
+	if err != nil {
+		return err
+	}
+
+	return nil
 }
 
-func handleEchoRequest(req EchoRequest) {
-	slog.Info("Got echo request", "EchoByte", req.EchoByte)
+func handleEchoRequest(echoRequestMsg Message) error {
+	echoRequest := echoRequestMsg.(EchoRequest)
+	slog.Info("Got echo request", "EchoByte", echoRequest.EchoByte)
+
+	return nil
 }
 
 func writeMessage(conn net.Conn, msg Message) error {
@@ -157,13 +181,7 @@ func encodeMessage(msg Message) ([]byte, error) {
 	var msgJsonBytes []byte
 	var err error
 
-	switch msg.ID() {
-	case EchoRequestID:
-		msgJsonBytes, err = json.Marshal(msg.(EchoRequest))
-	default:
-		err = errors.New("unknown message type")
-	}
-
+	msgJsonBytes, err = json.Marshal(msg)
 	if err != nil {
 		return nil, err
 	}