Commit 9295a3eb authored by Eric's avatar Eric

fix commmand parse

parents 2e4496e2 50388696
package proxy package proxy
import ( import (
"io" "bytes"
"github.com/jumpserver/koko/pkg/logger"
"github.com/jumpserver/koko/pkg/utils"
"regexp" "regexp"
"strings" "strings"
"sync" "sync"
"github.com/jumpserver/koko/pkg/logger"
"github.com/jumpserver/koko/pkg/utils"
) )
var ps1Pattern = regexp.MustCompile(`^\[?.*@.*\]?[\\$#]\s|mysql>\s`) var ps1Pattern = regexp.MustCompile(`^\[?.*@.*\]?[\\$#]\s|mysql>\s`)
...@@ -21,91 +20,29 @@ func NewCmdParser(sid, name string) *CmdParser { ...@@ -21,91 +20,29 @@ func NewCmdParser(sid, name string) *CmdParser {
type CmdParser struct { type CmdParser struct {
id string id string
name string name string
buf bytes.Buffer
term *utils.Terminal
reader io.ReadCloser
writer io.WriteCloser
currentLines []string
lock *sync.Mutex lock *sync.Mutex
maxLength int maxLength int
currentLength int currentLength int
closed chan struct{}
} }
func (cp *CmdParser) WriteData(p []byte) (int, error) { func (cp *CmdParser) WriteData(p []byte) (int, error) {
select { cp.lock.Lock()
case <-cp.closed: defer cp.lock.Unlock()
return 0, io.EOF if cp.buf.Len() >= 1024 {
default: return 0, nil
}
return cp.writer.Write(p)
}
func (cp *CmdParser) Write(p []byte) (int, error) {
select {
case <-cp.closed:
return 0, io.EOF
default:
}
return len(p), nil
}
func (cp *CmdParser) Read(p []byte) (int, error) {
select {
case <-cp.closed:
return 0, io.EOF
default:
} }
return cp.reader.Read(p) return cp.buf.Write(p)
} }
func (cp *CmdParser) Close() error { func (cp *CmdParser) Close() error {
select { logger.Infof("session ID: %s, parser name: %s", cp.id, cp.name)
case <-cp.closed: return nil
return nil
default:
close(cp.closed)
}
_ = cp.reader.Close()
return cp.writer.Close()
} }
func (cp *CmdParser) initial() { func (cp *CmdParser) initial() {
cp.reader, cp.writer = io.Pipe()
cp.currentLines = make([]string, 0)
cp.lock = new(sync.Mutex) cp.lock = new(sync.Mutex)
cp.maxLength = 1024
cp.currentLength = 0
cp.closed = make(chan struct{})
cp.term = utils.NewTerminal(cp, "")
cp.term.SetEcho(false)
go func() {
logger.Infof("Session %s: %s start", cp.id, cp.name)
defer logger.Infof("Session %s: %s close", cp.id, cp.name)
loop:
for {
line, err := cp.term.ReadLine()
if err != nil {
select {
case <-cp.closed:
logger.Debugf("Session %s %s term err: %s break loop", cp.id, cp.name, err)
break loop
default:
}
logger.Debugf("Session %s %s term err: %s,loop continue", cp.id, cp.name, err)
goto loop
}
cp.lock.Lock()
cp.currentLength += len(line)
if cp.currentLength < cp.maxLength {
cp.currentLines = append(cp.currentLines, line)
}
cp.lock.Unlock()
}
}()
} }
func (cp *CmdParser) parsePS1(s string) string { func (cp *CmdParser) parsePS1(s string) string {
...@@ -114,16 +51,11 @@ func (cp *CmdParser) parsePS1(s string) string { ...@@ -114,16 +51,11 @@ func (cp *CmdParser) parsePS1(s string) string {
// Parse 解析命令或输出 // Parse 解析命令或输出
func (cp *CmdParser) Parse() string { func (cp *CmdParser) Parse() string {
select {
case <-cp.closed:
default:
cp.writer.Write([]byte("\r"))
}
cp.lock.Lock() cp.lock.Lock()
defer cp.lock.Unlock() defer cp.lock.Unlock()
output := strings.TrimSpace(strings.Join(cp.currentLines, "\r\n")) lines := utils.ParseTerminalData(cp.buf.Bytes())
output := strings.TrimSpace(strings.Join(lines, "\r\n"))
output = cp.parsePS1(output) output = cp.parsePS1(output)
cp.currentLines = make([]string, 0) cp.buf.Reset()
cp.currentLength = 0
return output return output
} }
// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package utils package utils
import ( import (
"bytes" "bytes"
"fmt"
"unicode/utf8" "unicode/utf8"
"github.com/jumpserver/koko/pkg/logger"
) )
func ParseTerminalData(p []byte) (lines []string, ok bool) {
c := bytes.NewReader(p)
pasteActive := false
ok = true
var line []rune
var pos int
var remainder []byte
var inBuf [256]byte
for {
rest := remainder
lineOk := false
for !lineOk {
var key rune
key, rest = bytesToKey(rest, pasteActive)
if key == utf8.RuneError {
break
}
if !pasteActive {
if key == keyPasteStart {
pasteActive = true
if len(line) == 0 {
}
continue
}
} else if key == keyPasteEnd {
pasteActive = false
continue
}
switch key { type terminalParser struct {
case keyBackspace:
if pos == 0 {
continue
}
line, pos = EraseNPreviousChars(1, pos, line)
case keyAltLeft:
// move left by a word.
pos -= CountToLeftWord(pos, line)
case keyAltRight:
// move right by a word.
pos += CountToRightWord(pos, line)
case keyLeft:
if pos == 0 {
continue
}
pos--
case keyRight:
if pos == len(line) {
continue
}
pos++
case keyHome:
if pos == 0 {
continue
}
pos = 0
case keyEnd:
if pos == len(line) {
continue
}
pos = len(line)
case keyUp:
line = []rune{}
pos = 0
ok = false
case keyDown:
line = []rune{}
pos = 0
ok = false
case keyEnter:
lines = append(lines, string(line))
line = line[:0]
pos = 0
lineOk = true
case keyDeleteWord:
// Delete zero or more spaces and then one or more characters.
line, pos = EraseNPreviousChars(CountToLeftWord(pos, line), pos, line)
case keyDeleteLine:
line = line[:pos]
case keyCtrlD:
// Erase the character under the current position.
// The EOF case when the line is empty is handled in
// readLine().
if pos < len(line) {
pos++
line, pos = EraseNPreviousChars(1, pos, line)
}
case keyCtrlU:
line = line[:0]
case keyClearScreen:
default:
if !isPrintable(key) {
fmt.Println("could not printable: ", []byte(string(key)), " ", key)
ok = false
continue
}
line, pos = AddKeyToLine(key, pos, line)
}
} // line is the current line being entered.
if len(rest) > 0 { line []rune
n := copy(inBuf[:], rest) // pos is the logical position of the cursor in line
remainder = inBuf[:n] pos int
} else { // pasteActive is true iff there is a bracketed paste operation in
remainder = nil // progress.
} pasteActive bool
// remainder is a slice at the beginning of t.inBuf // maxLine is the greatest value of cursorY so far.
// containing a partial key sequence maxLine int
readBuf := inBuf[len(remainder):]
var n int // remainder contains the remainder of any partial key sequences after
n, err := c.Read(readBuf) // a read. It aliases into inBuf.
if err != nil { remainder []byte
if len(line) > 0 { inBuf [256]byte
lines = append(lines, string(line))
} else if len(rest) > 0 {
lines = append(lines, string(rest))
}
return // history contains previously entered commands so that they can be
} // accessed with the up and down keys.
remainder = inBuf[:n+len(remainder)] history stRingBuffer
} // historyIndex stores the currently accessed history entry, where zero
// means the immediately previous entry.
historyIndex int
// When navigating up and down the history it's possible to return to
// the incomplete, initial line. That value is stored in
// historyPending.
historyPending string
}
func (t *terminalParser) setLine(newLine []rune, newPos int) {
t.line = newLine
t.pos = newPos
} }
func EraseNPreviousChars(n, cPos int, line []rune) ([]rune, int) { func (t *terminalParser) eraseNPreviousChars(n int) {
if n == 0 { if n == 0 {
return line, cPos return
} }
if cPos < n {
n = cPos if t.pos < n {
n = t.pos
} }
cPos -= n t.pos -= n
copy(line[cPos:], line[n+cPos:])
return line[:len(line)-n], cPos copy(t.line[t.pos:], t.line[n+t.pos:])
t.line = t.line[:len(t.line)-n]
} }
func CountToLeftWord(currentPos int, line []rune) int { // countToLeftWord returns then number of characters from the cursor to the
if currentPos == 0 { // start of the previous word.
func (t *terminalParser) countToLeftWord() int {
if t.pos == 0 {
return 0 return 0
} }
pos := currentPos - 1 pos := t.pos - 1
for pos > 0 { for pos > 0 {
if line[pos] != ' ' { if t.line[pos] != ' ' {
break break
} }
pos-- pos--
} }
for pos > 0 { for pos > 0 {
if line[pos] == ' ' { if t.line[pos] == ' ' {
pos++ pos++
break break
} }
pos-- pos--
} }
return currentPos - pos return t.pos - pos
} }
func CountToRightWord(currentPos int, line []rune) int { // countToRightWord returns then number of characters from the cursor to the
pos := currentPos // start of the next word.
for pos < len(line) { func (t *terminalParser) countToRightWord() int {
if line[pos] == ' ' { pos := t.pos
for pos < len(t.line) {
if t.line[pos] == ' ' {
break break
} }
pos++ pos++
} }
for pos < len(line) { for pos < len(t.line) {
if line[pos] != ' ' { if t.line[pos] != ' ' {
break break
} }
pos++ pos++
} }
return pos - currentPos return pos - t.pos
}
// handleKey processes the given key and, optionally, returns a line of text
// that the user has entered.
func (t *terminalParser) handleKey(key rune) (line string, ok bool) {
if t.pasteActive && key != keyEnter {
t.addKeyToLine(key)
return
}
switch key {
case keyBackspace:
if t.pos == 0 {
return
}
t.eraseNPreviousChars(1)
case keyAltLeft:
// move left by a word.
t.pos -= t.countToLeftWord()
case keyAltRight:
// move right by a word.
t.pos += t.countToRightWord()
case keyLeft:
if t.pos == 0 {
return
}
t.pos--
case keyRight:
if t.pos == len(t.line) {
return
}
t.pos++
case keyHome:
if t.pos == 0 {
return
}
t.pos = 0
case keyEnd:
if t.pos == len(t.line) {
return
}
t.pos = len(t.line)
case keyUp:
entry, ok := t.history.NthPreviousEntry(t.historyIndex + 1)
if !ok {
return "", false
}
if t.historyIndex == -1 {
t.historyPending = string(t.line)
}
t.historyIndex++
runes := []rune(entry)
t.setLine(runes, len(runes))
case keyDown:
switch t.historyIndex {
case -1:
return
case 0:
runes := []rune(t.historyPending)
t.setLine(runes, len(runes))
t.historyIndex--
default:
entry, ok := t.history.NthPreviousEntry(t.historyIndex - 1)
if ok {
t.historyIndex--
runes := []rune(entry)
t.setLine(runes, len(runes))
}
}
case keyEnter:
line = string(t.line)
ok = true
t.line = t.line[:0]
t.pos = 0
t.maxLine = 0
case keyDeleteWord:
// Delete zero or more spaces and then one or more characters.
t.eraseNPreviousChars(t.countToLeftWord())
case keyDeleteLine:
t.line = t.line[:t.pos]
case keyCtrlD:
// Erase the character under the current position.
// The EOF case when the line is empty is handled in
// readLine().
if t.pos < len(t.line) {
t.pos++
t.eraseNPreviousChars(1)
}
case keyCtrlU:
t.eraseNPreviousChars(t.pos)
case keyClearScreen:
// Erases the screen and moves the cursor to the home position.
t.setLine(t.line, t.pos)
default:
if !isPrintable(key) {
return
}
if len(t.line) == maxLineLength {
return
}
t.addKeyToLine(key)
}
return
}
// addKeyToLine inserts the given key at the current position in the current
// line.
func (t *terminalParser) addKeyToLine(key rune) {
if len(t.line) == cap(t.line) {
newLine := make([]rune, len(t.line), 2*(1+len(t.line)))
copy(newLine, t.line)
t.line = newLine
}
t.line = t.line[:len(t.line)+1]
copy(t.line[t.pos+1:], t.line[t.pos:])
t.line[t.pos] = key
t.pos++
}
func (t *terminalParser) parseLines(p []byte) (lines []string) {
var err error
lines = make([]string, 0, 3)
lineIsPasted := t.pasteActive
reader := bytes.NewBuffer(p)
for {
rest := t.remainder
line := ""
lineOk := false
for !lineOk {
var key rune
key, rest = bytesToKey(rest, t.pasteActive)
if key == utf8.RuneError {
break
}
if !t.pasteActive {
if key == keyCtrlD {
if len(t.line) == 0 {
// as key has already handled, we need update remainder data,
t.remainder = rest
return lines
}
}
if key == keyPasteStart {
t.pasteActive = true
if len(t.line) == 0 {
lineIsPasted = true
}
continue
}
} else if key == keyPasteEnd {
t.pasteActive = false
continue
}
if !t.pasteActive {
lineIsPasted = false
}
line, lineOk = t.handleKey(key)
}
if len(rest) > 0 {
n := copy(t.inBuf[:], rest)
t.remainder = t.inBuf[:n]
} else {
t.remainder = nil
}
if lineOk {
if lineIsPasted {
err = ErrPasteIndicator
}
lines = append(lines, line)
}
// t.remainder is a slice at the beginning of t.inBuf
// containing a partial key sequence
readBuf := t.inBuf[len(t.remainder):]
var n int
n, err = reader.Read(readBuf)
if err != nil && n == 0 {
if len(t.line) > 0 {
lines = append(lines, string(t.line))
}
if len(t.remainder) > 0{
continue
}
return
} else if err == nil && n == 0 {
if len(t.remainder) == len(t.inBuf) {
logger.Errorf("~~ 发生卡顿问题 ~~")
t.remainder = t.remainder[1:]
continue
}
}
t.remainder = t.inBuf[:n+len(t.remainder)]
}
} }
func AddKeyToLine(key rune, pos int, line []rune) ([]rune, int) { func ParseTerminalData(p []byte) (lines []string) {
if len(line) == cap(line) { t := terminalParser{
newLine := make([]rune, len(line), 2*(1+len(line))) historyIndex: -1,
copy(newLine, line)
line = newLine
} }
line = line[:len(line)+1] return t.parseLines(p)
copy(line[pos+1:], line[pos:])
line[pos] = key
pos++
return line, pos
} }
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment