update func

errors.go

@@ -16,7 +16,9 @@
 
 package jrpc
 
-import "fmt"
+import (
+	"fmt"
+)
 
 // HTTPError is returned by client operations when the HTTP status code of the
 // response is not a 2xx status.
@@ -45,6 +47,33 @@ type DataError interface {
 	ErrorData() any // returns the error data
 }
 
+type JrpcErr struct {
+	Data any
+}
+
+func (j *JrpcErr) ErrorData() any {
+	return j.Data
+}
+
+func (j *JrpcErr) Error() string {
+	return "Jrpc Error"
+}
+
+func (j *JrpcErr) ErrorCode() int {
+	return jrpcErrorCode
+}
+
+func WrapJrpcErr(err error) error {
+	if err == nil {
+		return nil
+	}
+	return fmt.Errorf("%w: %w", &JrpcErr{}, err)
+}
+
+func MakeJrpcErr(s string) error {
+	return fmt.Errorf("%w: %s", &JrpcErr{}, s)
+}
+
 // Error types defined below are the built-in JSON-RPC errors.
 
 var (
@@ -60,6 +89,8 @@ const defaultErrorCode = -32000
 
 const applicationErrorCode = -32080
 
+const jrpcErrorCode = -42000
+
 type methodNotFoundError struct{ method string }
 
 func (e *methodNotFoundError) ErrorCode() int { return -32601 }

websocket_server.go

@@ -22,11 +22,13 @@ func isWebsocket(r *http.Request) bool {
 		strings.Contains(strings.ToLower(r.Header.Get("Connection")), "upgrade")
 }
 
-func (s *Server) ServeHTTPWithWss(cb func(r *http.Request)) http.Handler {
+func (s *Server) ServeHTTPWithWss(cb func(w http.ResponseWriter, r *http.Request) bool) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if isWebsocket(r) {
 			if cb != nil {
-				cb(r)
+				if cb(w, r) {
+					return
+				}
 			}
 			s.WebsocketHandler([]string{"*"}).ServeHTTP(w, r)
 			return