From b0180761d1ee3b7d8f40e6025cf017fb124fe6d1 Mon Sep 17 00:00:00 2001
From: AmarnathCJD <amarnathcharichilil@gmail.com>
Date: Thu, 19 Dec 2024 01:21:23 +0530
Subject: [PATCH] Add errorHandler to MTProto and ClientConfig for improved
 error handling

---
 mtproto.go         | 16 +++++++++++++---
 telegram/client.go |  3 +++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/mtproto.go b/mtproto.go
index eb9bfaaf..46dd4f53 100755
--- a/mtproto.go
+++ b/mtproto.go
@@ -80,6 +80,7 @@ type MTProto struct {
 
 	serverRequestHandlers []func(i any) bool
 	floodHandler          func(err error) bool
+	errorHandler          func(err error)
 	exported              bool
 }
 
@@ -89,7 +90,9 @@ type Config struct {
 	SessionStorage session.SessionLoader
 	MemorySession  bool
 	AppID          int32
-	FloodHandler   func(err error) bool
+
+	FloodHandler func(err error) bool
+	ErrorHandler func(err error)
 
 	ServerHost string
 	PublicKey  *rsa.PublicKey
@@ -142,6 +145,7 @@ func NewMTProto(c Config) (*MTProto, error) {
 		appID:                 c.AppID,
 		proxy:                 c.Proxy,
 		floodHandler:          func(err error) bool { return false },
+		errorHandler:          func(err error) {},
 		mode:                  parseTransportMode(c.Mode),
 		IpV6:                  c.Ipv6,
 	}
@@ -164,6 +168,10 @@ func NewMTProto(c Config) (*MTProto, error) {
 		mtproto.floodHandler = c.FloodHandler
 	}
 
+	if c.ErrorHandler != nil {
+		mtproto.errorHandler = c.ErrorHandler
+	}
+
 	return mtproto, nil
 }
 
@@ -409,8 +417,8 @@ func (m *MTProto) makeRequest(data tl.Object, expectedTypes ...reflect.Type) (an
 
 	resp, _, err := m.sendPacket(data, expectedTypes...)
 	if err != nil {
-		if strings.Contains(err.Error(), "use of closed network connection") || strings.Contains(err.Error(), "transport is closed") || strings.Contains(err.Error(), "connection was forcibly closed") {
-			m.Logger.Info("connection closed due to broken tcp, reconnecting to [" + m.Addr + "]" + " - <Tcp> ...")
+		if strings.Contains(err.Error(), "use of closed network connection") || strings.Contains(err.Error(), "transport is closed") || strings.Contains(err.Error(), "connection was forcibly closed") || strings.Contains(err.Error(), "connection reset by peer") || strings.Contains(err.Error(), "broken pipe") {
+			m.Logger.Info(fmt.Sprintf("connection closed due to broken tcp, reconnecting to [%s] - <%s> ...", m.Addr, utils.Vtcp(m.IpV6)))
 			err = m.Reconnect(false)
 			if err != nil {
 				return nil, errors.Wrap(err, "reconnecting")
@@ -429,6 +437,8 @@ func (m *MTProto) makeRequest(data tl.Object, expectedTypes ...reflect.Type) (an
 				return m.makeRequest(data, expectedTypes...)
 			}
 		}
+		m.errorHandler(RpcErrorToNative(r))
+
 		return nil, RpcErrorToNative(r)
 
 	case *errorSessionConfigsChanged:
diff --git a/telegram/client.go b/telegram/client.go
index e9d82fff..98971c80 100644
--- a/telegram/client.go
+++ b/telegram/client.go
@@ -85,6 +85,7 @@ type ClientConfig struct {
 	Cache         *CACHE
 	TransportMode string
 	FloodHandler  func(err error) bool
+	ErrorHandler  func(err error)
 }
 
 type Session struct {
@@ -172,6 +173,8 @@ func (c *Client) setupMTProto(config ClientConfig) error {
 		MemorySession: config.MemorySession,
 		Ipv6:          config.ForceIPv6,
 		CustomHost:    customHost,
+		FloodHandler:  config.FloodHandler,
+		ErrorHandler:  config.ErrorHandler,
 	})
 	if err != nil {
 		return errors.Wrap(err, "creating mtproto client")