diff --git a/cmd/go-judge/stream/file.go b/cmd/go-judge/stream/file.go index 22b9a5e..7b74a98 100644 --- a/cmd/go-judge/stream/file.go +++ b/cmd/go-judge/stream/file.go @@ -1,8 +1,8 @@ package stream import ( + "errors" "fmt" - "os" "github.com/criyle/go-judge/envexec" "github.com/criyle/go-judge/filestore" @@ -12,6 +12,8 @@ import ( var ( _ worker.CmdFile = &fileStreamIn{} _ worker.CmdFile = &fileStreamOut{} + + errNoTTY = errors.New("stream is not a tty") ) type fileStreamIn struct { @@ -30,15 +32,15 @@ func newFileStreamIn(index, fd int, hasTTY bool) *fileStreamIn { } } -func (f *fileStreamIn) GetTTY() *os.File { +func (f *fileStreamIn) SetSize(s *envexec.TerminalSize) error { if !f.hasTTY { - return nil + return errNoTTY } - return f.stream.WritePipe() + return f.stream.SetSize(s) } func (f *fileStreamIn) Write(b []byte) (int, error) { - return f.stream.WritePipe().Write(b) + return f.stream.Write(b) } func (f *fileStreamIn) EnvFile(fs filestore.FileStore) (envexec.File, error) { @@ -68,7 +70,7 @@ func newFileStreamOut(index, fd int) *fileStreamOut { } func (f *fileStreamOut) Read(b []byte) (int, error) { - return f.stream.ReadPipe().Read(b) + return f.stream.Read(b) } func (f *fileStreamOut) EnvFile(fs filestore.FileStore) (envexec.File, error) { diff --git a/cmd/go-judge/stream/stream.go b/cmd/go-judge/stream/stream.go index 77b7c93..e790fa0 100644 --- a/cmd/go-judge/stream/stream.go +++ b/cmd/go-judge/stream/stream.go @@ -7,6 +7,7 @@ import ( "io" "github.com/criyle/go-judge/cmd/go-judge/model" + "github.com/criyle/go-judge/envexec" "github.com/criyle/go-judge/worker" "go.uber.org/zap" "golang.org/x/sync/errgroup" @@ -242,11 +243,12 @@ func streamInput(ctx context.Context, s Stream, si []*fileStreamIn, execCancel f if !ok { return fmt.Errorf("input does not exist: %d/%d", in.Resize.Index, in.Resize.Fd) } - tty := f.GetTTY() - if tty == nil { - return fmt.Errorf("resize input does not have tty: %d/%d", in.Resize.Index, in.Resize.Fd) - } - if err = setWinsize(tty, in.Resize); err != nil { + if err = f.SetSize(&envexec.TerminalSize{ + Cols: uint16(in.Resize.Cols), + Rows: uint16(in.Resize.Rows), + X: uint16(in.Resize.X), + Y: uint16(in.Resize.Y), + }); err != nil { return fmt.Errorf("resize %d/%d: %w", in.Resize.Index, in.Resize.Fd, err) } diff --git a/cmd/go-judge/stream/util_other.go b/cmd/go-judge/stream/util_other.go deleted file mode 100644 index be1023f..0000000 --- a/cmd/go-judge/stream/util_other.go +++ /dev/null @@ -1,11 +0,0 @@ -//go:build !linux && !darwin - -package stream - -import ( - "os" -) - -func setWinsize(f *os.File, i *ResizeRequest) error { - return nil -} diff --git a/cmd/go-judge/stream/util_unix.go b/cmd/go-judge/stream/util_unix.go deleted file mode 100644 index 166bb22..0000000 --- a/cmd/go-judge/stream/util_unix.go +++ /dev/null @@ -1,19 +0,0 @@ -//go:build linux || darwin - -package stream - -import ( - "os" - - "github.com/creack/pty" -) - -func setWinsize(f *os.File, i *ResizeRequest) error { - winSize := &pty.Winsize{ - Rows: uint16(i.Rows), - Cols: uint16(i.Cols), - X: uint16(i.X), - Y: uint16(i.Y), - } - return pty.Setsize(f, winSize) -} diff --git a/envexec/file.go b/envexec/file.go index 38329ee..0b5a663 100644 --- a/envexec/file.go +++ b/envexec/file.go @@ -5,6 +5,8 @@ import ( "io" "os" "sync" + + "github.com/creack/pty" ) var ( @@ -42,43 +44,82 @@ func NewFileReader(r io.Reader) File { // responsibility to close the WritePipe type FileStreamIn interface { File - Done() <-chan struct{} - WritePipe() *os.File - Close() error + io.WriteCloser + SetSize(*TerminalSize) error +} + +// TerminalSize controls the size of the terminal if TTY is enabled +type TerminalSize struct { + Rows uint16 // ws_row: Number of rows (in cells). + Cols uint16 // ws_col: Number of columns (in cells). + X uint16 // ws_xpixel: Width in pixels. + Y uint16 // ws_ypixel: Height in pixels. } type fileStreamIn struct { - done chan struct{} - w *sharedFile + started chan struct{} + closed chan struct{} + w *sharedFile } func NewFileStreamIn() FileStreamIn { return &fileStreamIn{ - done: make(chan struct{}), + started: make(chan struct{}), + closed: make(chan struct{}), } } func (*fileStreamIn) isFile() {} func (f *fileStreamIn) start(w *sharedFile) { - f.w = w - close(f.done) + select { + case <-f.closed: + w.Close() + + default: + f.w = w + close(f.started) + } } -func (f *fileStreamIn) Done() <-chan struct{} { - return f.done +func (f *fileStreamIn) Write(b []byte) (int, error) { + select { + case <-f.started: + return f.w.f.Write(b) + + case <-f.closed: + return 0, io.EOF + } } -func (f *fileStreamIn) WritePipe() *os.File { - <-f.done - return f.w.f +func (f *fileStreamIn) SetSize(s *TerminalSize) error { + select { + case <-f.started: + return pty.Setsize(f.w.f, &pty.Winsize{ + Rows: s.Rows, + Cols: s.Cols, + X: s.X, + Y: s.Y, + }) + + case <-f.closed: + return io.ErrClosedPipe + } } func (f *fileStreamIn) Close() error { - if f.w != nil { + select { + case <-f.started: + close(f.closed) return f.w.Close() + + case <-f.closed: + return io.ErrClosedPipe + + default: + close(f.closed) + return nil } - return nil } // FileStreamOut represent a out streaming pipe and the streamer is able to read @@ -86,43 +127,58 @@ func (f *fileStreamIn) Close() error { // responsibility to close the ReadPipe type FileStreamOut interface { File - Done() <-chan struct{} - ReadPipe() *os.File - Close() error + io.ReadCloser } type fileStreamOut struct { - done chan struct{} - r *sharedFile + started chan struct{} + closed chan struct{} + r *sharedFile } func NewFileStreamOut() FileStreamOut { return &fileStreamOut{ - done: make(chan struct{}), + started: make(chan struct{}), + closed: make(chan struct{}), } } func (*fileStreamOut) isFile() {} func (f *fileStreamOut) start(r *sharedFile) { - f.r = r - close(f.done) + select { + case <-f.closed: + r.Close() + + default: + f.r = r + close(f.started) + } } -func (f *fileStreamOut) Done() <-chan struct{} { - return f.done -} +func (f *fileStreamOut) Read(p []byte) (n int, err error) { + select { + case <-f.started: + return f.r.f.Read(p) -func (f *fileStreamOut) ReadPipe() *os.File { - <-f.done - return f.r.f + case <-f.closed: + return 0, io.EOF + } } func (f *fileStreamOut) Close() error { - if f.r != nil { + select { + case <-f.started: + close(f.closed) return f.r.Close() + + case <-f.closed: + return io.ErrClosedPipe + + default: + close(f.closed) + return nil } - return nil } // FileInput represent file input which will be opened in read-only mode @@ -204,7 +260,7 @@ type sharedFile struct { count int } -func newShreadFile(f *os.File) *sharedFile { +func newSharedFile(f *os.File) *sharedFile { return &sharedFile{f: f, count: 0} } diff --git a/envexec/file_prepare.go b/envexec/file_prepare.go index e64d8e7..df772fb 100644 --- a/envexec/file_prepare.go +++ b/envexec/file_prepare.go @@ -31,7 +31,7 @@ func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (f []*os.File err = fmt.Errorf("failed to open tty %v", err) return nil, nil, err } - sf := newShreadFile(fPty) + sf := newSharedFile(fPty) files := make([]*os.File, count) pipeToCollect := make([]pipeCollector, 0) @@ -145,6 +145,34 @@ func prepareCmdFdNoTty(c *Cmd, count int, newFileStore NewStoreFile) (f []*os.Fi }() // record the same file to avoid multiple file open cf := make(map[string]*os.File) + prepareFileCollector := func(t *FileCollector) (*os.File, error) { + if f, ok := cf[t.Name]; ok { + return f, nil + } + + if t.Pipe { + b, err := newPipeBuffer(t.Limit, newFileStore) + if err != nil { + return nil, fmt.Errorf("pipe: create: %w", err) + } + cf[t.Name] = b.W + pipeToCollect = append(pipeToCollect, pipeCollector{b.Done, b.Buffer, t.Limit, t.Name, true}) + return b.W, nil + } else { + f, err := c.Environment.Open(t.Name, os.O_CREATE|os.O_WRONLY, 0777) + if err != nil { + return nil, fmt.Errorf("container: create file %v: %w", t.Name, err) + } + buffer, err := c.Environment.Open(t.Name, os.O_RDWR, 0777) + if err != nil { + f.Close() + return nil, fmt.Errorf("container: open created file %v: %w", t.Name, err) + } + cf[t.Name] = f + pipeToCollect = append(pipeToCollect, pipeCollector{closedChan, buffer, t.Limit, t.Name, false}) + return f, nil + } + } for j, t := range c.Files { switch t := t.(type) { @@ -167,35 +195,11 @@ func prepareCmdFdNoTty(c *Cmd, count int, newFileStore NewStoreFile) (f []*os.Fi files[j] = f case *FileCollector: - if f, ok := cf[t.Name]; ok { - files[j] = f - break - } - - if t.Pipe { - b, err := newPipeBuffer(t.Limit, newFileStore) - if err != nil { - return nil, nil, fmt.Errorf("pipe: create: %w", err) - } - cf[t.Name] = b.W - - files[j] = b.W - pipeToCollect = append(pipeToCollect, pipeCollector{b.Done, b.Buffer, t.Limit, t.Name, true}) - } else { - f, err := c.Environment.Open(t.Name, os.O_CREATE|os.O_WRONLY, 0777) - if err != nil { - return nil, nil, fmt.Errorf("container: create file %v: %w", t.Name, err) - } - cf[t.Name] = f - - buffer, err := c.Environment.Open(t.Name, os.O_RDWR, 0777) - if err != nil { - return nil, nil, fmt.Errorf("container: open created file %v: %w", t.Name, err) - } - - files[j] = f - pipeToCollect = append(pipeToCollect, pipeCollector{closedChan, buffer, t.Limit, t.Name, false}) + f, err := prepareFileCollector(t) + if err != nil { + return nil, nil, err } + files[j] = f case *FileWriter: _, w, err := newPipe(t.Writer, t.Limit)