From fad55dce28a7a244c7366d1319beacefb69fab8f Mon Sep 17 00:00:00 2001
From: naison <895703375@qq.com>
Date: Wed, 18 Dec 2024 20:50:47 +0800
Subject: [PATCH] feat(log): log trace if panic (#405)

---
 pkg/core/gvisortcphandler.go       |  2 ++
 pkg/core/gvisorudpforwarder.go     |  3 +++
 pkg/core/gvisorudphandler.go       |  3 +++
 pkg/core/tunhandler.go             |  6 +++++
 pkg/core/tunhandlerclient.go       |  4 ++++
 pkg/daemon/daemon.go               |  8 ++++---
 pkg/daemon/rpc/panicinterceptor.go | 37 ++++++++++++++++++++++++++++++
 pkg/dns/dns.go                     |  2 ++
 pkg/util/grpc.go                   | 11 ++++++++-
 9 files changed, 72 insertions(+), 4 deletions(-)
 create mode 100644 pkg/daemon/rpc/panicinterceptor.go

diff --git a/pkg/core/gvisortcphandler.go b/pkg/core/gvisortcphandler.go
index 51b9e8c0..915d5e77 100644
--- a/pkg/core/gvisortcphandler.go
+++ b/pkg/core/gvisortcphandler.go
@@ -40,10 +40,12 @@ func (h *gvisorTCPHandler) handle(ctx context.Context, tcpConn net.Conn) {
 	endpoint := channel.New(tcp.DefaultReceiveBufferSize, uint32(config.DefaultMTU), tcpip.GetRandMacAddr())
 	errChan := make(chan error, 2)
 	go func() {
+		defer util.HandleCrash()
 		h.readFromTCPConnWriteToEndpoint(ctx, tcpConn, endpoint)
 		util.SafeClose(errChan)
 	}()
 	go func() {
+		defer util.HandleCrash()
 		h.readFromEndpointWriteToTCPConn(ctx, tcpConn, endpoint)
 		util.SafeClose(errChan)
 	}()
diff --git a/pkg/core/gvisorudpforwarder.go b/pkg/core/gvisorudpforwarder.go
index 4e956dd9..1225d5a4 100644
--- a/pkg/core/gvisorudpforwarder.go
+++ b/pkg/core/gvisorudpforwarder.go
@@ -14,6 +14,7 @@ import (
 	"gvisor.dev/gvisor/pkg/waiter"
 
 	"github.com/wencaiwulue/kubevpn/v2/pkg/config"
+	"github.com/wencaiwulue/kubevpn/v2/pkg/util"
 )
 
 func UDPForwarder(s *stack.Stack, ctx context.Context) func(id stack.TransportEndpointID, pkt *stack.PacketBuffer) bool {
@@ -51,6 +52,7 @@ func UDPForwarder(s *stack.Stack, ctx context.Context) func(id stack.TransportEn
 			defer remote.Close()
 			errChan := make(chan error, 2)
 			go func() {
+				defer util.HandleCrash()
 				buf := config.LPool.Get().([]byte)[:]
 				defer config.LPool.Put(buf[:])
 
@@ -80,6 +82,7 @@ func UDPForwarder(s *stack.Stack, ctx context.Context) func(id stack.TransportEn
 				errChan <- err
 			}()
 			go func() {
+				defer util.HandleCrash()
 				buf := config.LPool.Get().([]byte)[:]
 				defer config.LPool.Put(buf[:])
 
diff --git a/pkg/core/gvisorudphandler.go b/pkg/core/gvisorudphandler.go
index e3d658fe..3a98d217 100644
--- a/pkg/core/gvisorudphandler.go
+++ b/pkg/core/gvisorudphandler.go
@@ -8,6 +8,7 @@ import (
 
 	log "github.com/sirupsen/logrus"
 	"github.com/wencaiwulue/kubevpn/v2/pkg/config"
+	"github.com/wencaiwulue/kubevpn/v2/pkg/util"
 )
 
 type gvisorUDPHandler struct{}
@@ -102,6 +103,7 @@ func handle(ctx context.Context, tcpConn net.Conn, udpConn *net.UDPConn) {
 	log.Debugf("[TUN-UDP] %s <-> %s", tcpConn.RemoteAddr(), udpConn.LocalAddr())
 	errChan := make(chan error, 2)
 	go func() {
+		defer util.HandleCrash()
 		buf := config.LPool.Get().([]byte)[:]
 		defer config.LPool.Put(buf[:])
 
@@ -146,6 +148,7 @@ func handle(ctx context.Context, tcpConn net.Conn, udpConn *net.UDPConn) {
 	}()
 
 	go func() {
+		defer util.HandleCrash()
 		buf := config.LPool.Get().([]byte)[:]
 		defer config.LPool.Put(buf[:])
 
diff --git a/pkg/core/tunhandler.go b/pkg/core/tunhandler.go
index b99f7c64..392cad15 100644
--- a/pkg/core/tunhandler.go
+++ b/pkg/core/tunhandler.go
@@ -95,6 +95,7 @@ type Device struct {
 }
 
 func (d *Device) readFromTun() {
+	defer util.HandleCrash()
 	for {
 		buf := config.LPool.Get().([]byte)[:]
 		n, err := d.tun.Read(buf[:])
@@ -128,6 +129,7 @@ func (d *Device) readFromTun() {
 }
 
 func (d *Device) writeToTun() {
+	defer util.HandleCrash()
 	for e := range d.tunOutbound {
 		_, err := d.tun.Write(e.data[:e.length])
 		config.LPool.Put(e.data[:])
@@ -277,6 +279,7 @@ func (p *Peer) sendErr(err error) {
 }
 
 func (p *Peer) readFromConn() {
+	defer util.HandleCrash()
 	for {
 		buf := config.LPool.Get().([]byte)[:]
 		n, from, err := p.conn.ReadFrom(buf[:])
@@ -312,6 +315,7 @@ func (p *Peer) readFromConn() {
 }
 
 func (p *Peer) readFromTCPConn() {
+	defer util.HandleCrash()
 	for packet := range TCPPacketChan {
 		src, dst, err := util.ParseIP(packet.Data)
 		if err != nil {
@@ -331,6 +335,7 @@ func (p *Peer) readFromTCPConn() {
 }
 
 func (p *Peer) routePeer() {
+	defer util.HandleCrash()
 	for e := range p.connInbound {
 		if routeToAddr := p.routeMapUDP.RouteTo(e.dst); routeToAddr != nil {
 			log.Debugf("[UDP] Find UDP route to dst: %s -> %s", e.dst, routeToAddr)
@@ -363,6 +368,7 @@ func (p *Peer) routePeer() {
 }
 
 func (p *Peer) routeTUN() {
+	defer util.HandleCrash()
 	for e := range p.tunInbound {
 		if addr := p.routeMapUDP.RouteTo(e.dst); addr != nil {
 			log.Debugf("[TUN] Find UDP route to dst: %s -> %s", e.dst, addr)
diff --git a/pkg/core/tunhandlerclient.go b/pkg/core/tunhandlerclient.go
index bcfb71f5..59b14cec 100644
--- a/pkg/core/tunhandlerclient.go
+++ b/pkg/core/tunhandlerclient.go
@@ -81,6 +81,7 @@ func transportTunClient(ctx context.Context, tunInbound <-chan *DataElem, tunOut
 	defer packetConn.Close()
 
 	go func() {
+		defer util.HandleCrash()
 		for e := range tunInbound {
 			if e.src.Equal(e.dst) {
 				util.SafeWrite(tunOutbound, e)
@@ -96,6 +97,7 @@ func transportTunClient(ctx context.Context, tunInbound <-chan *DataElem, tunOut
 	}()
 
 	go func() {
+		defer util.HandleCrash()
 		for {
 			buf := config.LPool.Get().([]byte)[:]
 			n, _, err := packetConn.ReadFrom(buf[:])
@@ -145,6 +147,7 @@ func (d *ClientDevice) SetTunInboundHandler(handler func(tunInbound <-chan *Data
 }
 
 func (d *ClientDevice) readFromTun() {
+	defer util.HandleCrash()
 	for {
 		buf := config.LPool.Get().([]byte)[:]
 		n, err := d.tun.Read(buf[:])
@@ -172,6 +175,7 @@ func (d *ClientDevice) readFromTun() {
 }
 
 func (d *ClientDevice) writeToTun() {
+	defer util.HandleCrash()
 	for e := range d.tunOutbound {
 		_, err := d.tun.Write(e.data[:e.length])
 		config.LPool.Put(e.data[:])
diff --git a/pkg/daemon/daemon.go b/pkg/daemon/daemon.go
index 2d957af0..9c5e0e18 100644
--- a/pkg/daemon/daemon.go
+++ b/pkg/daemon/daemon.go
@@ -50,12 +50,12 @@ func (o *SvrOption) Start(ctx context.Context) error {
 		LocalTime:  true,
 		Compress:   false,
 	}
-	
+
 	// for gssapi to lookup KDCs in DNS
 	// c.LibDefaults.DNSLookupKDC = true
 	// c.LibDefaults.DNSLookupRealm = true
 	net.DefaultResolver.PreferGo = true
-	
+
 	util.InitLoggerForServer(true)
 	log.SetOutput(l)
 	klog.SetOutput(l)
@@ -89,7 +89,9 @@ func (o *SvrOption) Start(ctx context.Context) error {
 		return err
 	}
 
-	svr := grpc.NewServer()
+	unaryPanicInterceptor := grpc.UnaryInterceptor(rpc.UnaryPanicHandler)
+	streamPanicInterceptor := grpc.StreamInterceptor(rpc.StreamPanicHandler)
+	svr := grpc.NewServer(unaryPanicInterceptor, streamPanicInterceptor)
 	cleanup, err := admin.Register(svr)
 	if err != nil {
 		log.Errorf("Failed to register admin: %v", err)
diff --git a/pkg/daemon/rpc/panicinterceptor.go b/pkg/daemon/rpc/panicinterceptor.go
new file mode 100644
index 00000000..3470a47c
--- /dev/null
+++ b/pkg/daemon/rpc/panicinterceptor.go
@@ -0,0 +1,37 @@
+package rpc
+
+import (
+	"fmt"
+	"runtime/debug"
+
+	"github.com/sirupsen/logrus"
+	"golang.org/x/net/context"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+var _ grpc.UnaryServerInterceptor = UnaryPanicHandler
+var _ grpc.StreamServerInterceptor = StreamPanicHandler
+
+func UnaryPanicHandler(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			str := fmt.Sprintf("Panic: `%s` %s", info.FullMethod, string(debug.Stack()))
+			err = status.Error(codes.Internal, str)
+			logrus.Panic(str)
+		}
+	}()
+	return handler(ctx, req)
+}
+
+func StreamPanicHandler(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) (err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			str := fmt.Sprintf("Panic: `%s` %s", info.FullMethod, string(debug.Stack()))
+			err = status.Error(codes.Internal, str)
+			logrus.Panic(str)
+		}
+	}()
+	return handler(srv, ss)
+}
diff --git a/pkg/dns/dns.go b/pkg/dns/dns.go
index f3724b60..cf35d014 100644
--- a/pkg/dns/dns.go
+++ b/pkg/dns/dns.go
@@ -27,6 +27,7 @@ import (
 	"tailscale.com/net/dns"
 
 	"github.com/wencaiwulue/kubevpn/v2/pkg/config"
+	"github.com/wencaiwulue/kubevpn/v2/pkg/util"
 )
 
 type Config struct {
@@ -73,6 +74,7 @@ func (c *Config) AddServiceNameToHosts(ctx context.Context, serviceInterface v13
 }
 
 func (c *Config) watchServiceToAddHosts(ctx context.Context, serviceInterface v13.ServiceInterface, hosts []Entry) {
+	defer util.HandleCrash()
 	ticker := time.NewTicker(time.Second * 15)
 	defer ticker.Stop()
 	immediate := make(chan struct{}, 1)
diff --git a/pkg/util/grpc.go b/pkg/util/grpc.go
index 88a69615..3ec8719a 100644
--- a/pkg/util/grpc.go
+++ b/pkg/util/grpc.go
@@ -6,7 +6,9 @@ import (
 	"fmt"
 	"io"
 	"os"
-
+	"runtime/debug"
+	
+	"github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
 )
 
@@ -76,3 +78,10 @@ func CopyAndConvertGRPCStream[I any, O any](r grpc.ClientStream, w grpc.ServerSt
 		}
 	}
 }
+
+func HandleCrash() {
+	if r := recover(); r != nil {
+		logrus.Panic(fmt.Sprintf("Panic: %s", string(debug.Stack())))
+		panic(r)
+	}
+}