refactor: remove unuse variable and channel

This commit is contained in:
2025-07-21 12:10:08 +07:00
parent 32f9765374
commit 66f1492466
3 changed files with 24 additions and 37 deletions

BIN
main.exe Normal file

Binary file not shown.

View File

@ -41,8 +41,7 @@ var (
type Session struct { type Session struct {
Connection *ssh.ServerConn Connection *ssh.ServerConn
ConnChannels []ssh.Channel ConnChannel ssh.Channel
GlobalRequest <-chan *ssh.Request
Listener net.Listener Listener net.Listener
TunnelType TunnelType TunnelType TunnelType
ForwardedPort uint16 ForwardedPort uint16
@ -94,8 +93,8 @@ func (s *Session) Close() {
s.Listener.Close() s.Listener.Close()
} }
for _, ch := range s.ConnChannels { if s.ConnChannel != nil {
ch.Close() s.ConnChannel.Close()
} }
if s.Connection != nil { if s.Connection != nil {
@ -109,17 +108,19 @@ func (s *Session) Close() {
close(s.Done) close(s.Done)
} }
func (s *Session) handleGlobalRequest() { func (s *Session) handleGlobalRequest(GlobalRequest <-chan *ssh.Request) {
ticker := time.NewTicker(1 * time.Second) ticker := time.NewTicker(1 * time.Second)
for { for {
select { select {
case req := <-s.GlobalRequest: case req := <-GlobalRequest:
ticker.Stop() ticker.Stop()
if req == nil { if req == nil {
return return
} }
if req.Type == "tcpip-forward" { if req.Type == "tcpip-forward" {
s.handleTCPIPForward(req) s.handleTCPIPForward(req)
} else if req.Type == "shell" || req.Type == "pty-req" || req.Type == "window-change" {
req.Reply(true, nil)
} else { } else {
req.Reply(false, nil) req.Reply(false, nil)
} }
@ -128,6 +129,7 @@ func (s *Session) handleGlobalRequest() {
case <-ticker.C: case <-ticker.C:
s.sendMessage(fmt.Sprintf("Please specify the forwarding tunnel. For example: 'ssh %s -p %s -R 443:localhost:8080' \r\n\n\n", utils.Getenv("domain"), utils.Getenv("port"))) s.sendMessage(fmt.Sprintf("Please specify the forwarding tunnel. For example: 'ssh %s -p %s -R 443:localhost:8080' \r\n\n\n", utils.Getenv("domain"), utils.Getenv("port")))
s.Close() s.Close()
return
} }
} }
} }
@ -162,7 +164,7 @@ func (s *Session) handleTCPIPForward(req *ssh.Request) {
} }
s.sendMessage("\033[H\033[2J") s.sendMessage("\033[H\033[2J")
showWelcomeMessage(s.ConnChannels[0]) showWelcomeMessage(s.ConnChannel)
s.Status = RUNNING s.Status = RUNNING
if portToBind == 80 || portToBind == 443 { if portToBind == 80 || portToBind == 443 {
@ -299,23 +301,24 @@ func (s *Session) waitForRunningStatus() {
} }
func (s *Session) sendMessage(message string) { func (s *Session) sendMessage(message string) {
if len(s.ConnChannels) > 0 { if s.ConnChannel != nil {
s.ConnChannels[0].Write([]byte(message)) s.ConnChannel.Write([]byte(message))
} }
} }
func (s *Session) HandleSessionChannel(newChannel ssh.NewChannel) { func (s *Session) HandleSessionChannel(newChannel ssh.NewChannel, initialRequest <-chan *ssh.Request) {
connection, requests, err := newChannel.Accept() connection, requests, err := newChannel.Accept()
if err != nil { if err != nil {
log.Printf("Could not accept channel: %s", err) log.Printf("Could not accept channel: %s", err)
return return
} }
s.ConnChannels = append(s.ConnChannels, connection) s.ConnChannel = connection
s.Status = RUNNING
go s.handleGlobalRequest(initialRequest)
go s.handleGlobalRequest(requests)
go s.handleUserInput(connection) go s.handleUserInput(connection)
go s.handleChannelRequests(connection, requests)
} }
func (s *Session) handleUserInput(connection ssh.Channel) { func (s *Session) handleUserInput(connection ssh.Channel) {
@ -493,7 +496,7 @@ func (s *Session) handleCommand(connection ssh.Channel, command string, inSlugEd
connection.Write([]byte("\r\nAvailable commands: /bye, /help, /clear, /slug")) connection.Write([]byte("\r\nAvailable commands: /bye, /help, /clear, /slug"))
case "/clear": case "/clear":
connection.Write([]byte("\033[H\033[2J")) connection.Write([]byte("\033[H\033[2J"))
showWelcomeMessage(s.ConnChannels[0]) showWelcomeMessage(s.ConnChannel)
domain := utils.Getenv("domain") domain := utils.Getenv("domain")
if s.TunnelType == HTTP { if s.TunnelType == HTTP {
protocol := "http" protocol := "http"
@ -522,20 +525,6 @@ func (s *Session) handleCommand(connection ssh.Channel, command string, inSlugEd
commandBuffer.Reset() commandBuffer.Reset()
} }
func (s *Session) handleChannelRequests(connection ssh.Channel, requests <-chan *ssh.Request) {
go s.handleGlobalRequest()
for req := range requests {
switch req.Type {
case "shell", "pty-req", "window-change":
req.Reply(true, nil)
default:
log.Println("Unknown request type:", req.Type)
req.Reply(false, nil)
}
}
}
func (s *Session) HandleForwardedConnection(conn UserConnection, sshConn *ssh.ServerConn) { func (s *Session) HandleForwardedConnection(conn UserConnection, sshConn *ssh.ServerConn) {
defer conn.Writer.Close() defer conn.Writer.Close()

View File

@ -15,19 +15,17 @@ const (
func New(conn *ssh.ServerConn, sshChannel <-chan ssh.NewChannel, req <-chan *ssh.Request) *Session { func New(conn *ssh.ServerConn, sshChannel <-chan ssh.NewChannel, req <-chan *ssh.Request) *Session {
session := &Session{ session := &Session{
Status: SETUP, Status: SETUP,
Slug: "", Slug: "",
ConnChannels: []ssh.Channel{}, ConnChannel: nil,
Connection: conn, Connection: conn,
GlobalRequest: req, TunnelType: UNKNOWN,
TunnelType: UNKNOWN, Done: make(chan bool),
SlugChannel: make(chan bool),
Done: make(chan bool),
} }
go func() { go func() {
for newChannel := range sshChannel { for newChannel := range sshChannel {
go session.HandleSessionChannel(newChannel) go session.HandleSessionChannel(newChannel, req)
} }
}() }()