ServerError type check before EOF string comparison

lib/eof.go

@@ -2,7 +2,9 @@ package lib
 
 import (
 	"errors"
+	"fmt"
 	"io"
+	"net/rpc"
 	"strings"
 
 	"github.com/hashicorp/yamux"
@@ -20,10 +22,17 @@ func IsErrEOF(err error) bool {
 
 	errStr := err.Error()
 	if strings.Contains(errStr, yamuxStreamClosed) ||
-		strings.Contains(errStr, yamuxSessionShutdown) ||
+		strings.Contains(errStr, yamuxSessionShutdown) {
-		strings.HasSuffix(errStr, io.EOF.Error()) {
 		return true
 	}
 
+	if srvErr, ok := err.(rpc.ServerError); ok {
+		return strings.HasSuffix(srvErr.Error(), fmt.Sprintf(": %s", io.EOF.Error()))
+	}
+
+	if srvErr, ok := errors.Unwrap(err).(rpc.ServerError); ok {
+		return strings.HasSuffix(srvErr.Error(), fmt.Sprintf(": %s", io.EOF.Error()))
+	}
+
 	return false
 }

lib/eof_test.go

@@ -3,6 +3,7 @@ package lib
 import (
 	"fmt"
 	"io"
+	"net/rpc"
 	"testing"
 
 	"github.com/hashicorp/yamux"
@@ -15,19 +16,16 @@ func TestErrIsEOF(t *testing.T) {
 		err  error
 	}{
 		{name: "EOF", err: io.EOF},
+		{name: "Wrapped EOF", err: fmt.Errorf("test: %w", io.EOF)},
 		{name: "yamuxStreamClosed", err: yamux.ErrStreamClosed},
 		{name: "yamuxSessionShutdown", err: yamux.ErrSessionShutdown},
+		{name: "ServerError(___: EOF)", err: rpc.ServerError(fmt.Sprintf("rpc error: %s", io.EOF.Error()))},
+		{name: "Wrapped ServerError(___: EOF)", err: fmt.Errorf("rpc error: %w", rpc.ServerError(fmt.Sprintf("rpc error: %s", io.EOF.Error())))},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			require.True(t, IsErrEOF(tt.err))
 		})
-		t.Run(fmt.Sprintf("Wrapped %s", tt.name), func(t *testing.T) {
-			require.True(t, IsErrEOF(fmt.Errorf("test: %w", tt.err)))
-		})
-		t.Run(fmt.Sprintf("String suffix is %s", tt.name), func(t *testing.T) {
-			require.True(t, IsErrEOF(fmt.Errorf("rpc error: %s", tt.err.Error())))
-		})
 	}
 }