Add copyOutDir to Cmd

This commit is contained in:
criyle 2020-03-07 05:45:37 -05:00
parent 51ad40a869
commit bb9c448601
5 changed files with 92 additions and 5 deletions

View File

@ -83,6 +83,8 @@ interface Cmd {
copyOut?: string[];
// similar to copyOut but stores file in executor service and returns fileId, later download through /file/:fileId
copyOutCached?: string[];
// specifies the directory to dump container /w content
copyOutDir: string
}
enum Status {
@ -150,7 +152,8 @@ Single (Require `apt install g++` inside the container):
}
},
"copyOut": ["stdout", "stderr"],
"copyOutCached": ["a.cc", "a"]
"copyOutCached": ["a.cc", "a"],
"copyOutDir": "1"
}]
}
```

View File

@ -24,6 +24,7 @@ type cmd struct {
CopyOut []string `json:"copyOut"`
CopyOutCached []string `json:"copyOutCached"`
CopyOutDir string `json:"copyOutDir"`
}
type pipeIndex struct {
@ -52,6 +53,6 @@ type response struct {
Error string `json:"error,omitempty"`
Time uint64 `json:"time"`
Memory uint64 `json:"memory"`
Files map[string]string `json:"files"`
FileIDs map[string]string `json:"fileIds"`
Files map[string]string `json:"files,omitempty"`
FileIDs map[string]string `json:"fileIds,omitempty"`
}

View File

@ -3,6 +3,7 @@ package main
import (
"context"
"fmt"
"path"
"sync"
"time"
@ -175,6 +176,10 @@ func prepareCmd(rc cmd) (*envexec.Cmd, map[string]bool, error) {
copyIn, err := prepareCopyIn(rc.CopyIn)
copyOutSet := make(map[string]bool)
// pipe default copyout
for k := range pipeFileName {
copyOutSet[k] = true
}
copyOut := make([]string, 0, len(rc.CopyOut)+len(rc.CopyOutCached))
for _, fn := range rc.CopyOut {
if !pipeFileName[fn] {
@ -185,6 +190,8 @@ func prepareCmd(rc cmd) (*envexec.Cmd, map[string]bool, error) {
for _, fn := range rc.CopyOutCached {
if !pipeFileName[fn] {
copyOut = append(copyOut, fn)
} else {
delete(copyOutSet, fn)
}
}
@ -201,6 +208,7 @@ func prepareCmd(rc cmd) (*envexec.Cmd, map[string]bool, error) {
ProcLimit: rc.ProcLimit,
CopyIn: copyIn,
CopyOut: copyOut,
CopyOutDir: path.Join(*dir, rc.CopyOutDir),
Waiter: w.Wait,
}, copyOutSet, nil
}

View File

@ -31,6 +31,9 @@ type Cmd struct {
// file names to copyout after exec
CopyOut []string
// CopyOutDir specifies a dir to dump all /w contnet
CopyOutDir string
// Waiter is called after cmd starts and it should return
// once time limit exceeded.
// return true to as TLE and false as normal exits (context finished)

View File

@ -4,6 +4,7 @@ import (
"io/ioutil"
"os"
"sync"
"syscall"
"github.com/criyle/go-judge/file"
"github.com/criyle/go-sandbox/container"
@ -13,7 +14,6 @@ import (
func copyOutAndCollect(m container.Environment, c *Cmd, ptc []pipeCollector) (map[string]file.File, error) {
// wait to complete
var wg sync.WaitGroup
wg.Add(len(ptc))
fc := make(chan file.File, len(ptc)+len(c.CopyOut))
errC := make(chan error, 1) // collect only 1 error
@ -34,10 +34,10 @@ func copyOutAndCollect(m container.Environment, c *Cmd, ptc []pipeCollector) (ma
// open all
cFiles, err = m.Open(openCmd)
wg.Add(len(cFiles))
}
// copy out
wg.Add(len(cFiles))
for i := range cFiles {
go func(cFile *os.File, n string) {
defer wg.Done()
@ -52,6 +52,7 @@ func copyOutAndCollect(m container.Environment, c *Cmd, ptc []pipeCollector) (ma
}
// collect pipe
wg.Add(len(ptc))
for _, p := range ptc {
go func(p pipeCollector) {
defer wg.Done()
@ -63,6 +64,51 @@ func copyOutAndCollect(m container.Environment, c *Cmd, ptc []pipeCollector) (ma
}(p)
}
// copy out dir
if c.CopyOutDir != "" {
wg.Add(1)
go func() {
defer wg.Done()
// open container /w
openCmd := []container.OpenCmd{{
Path: "/w",
Flag: os.O_RDONLY,
}}
dir, err := m.Open(openCmd)
if err != nil {
writeErrorC(errC, err)
return
}
defer dir[0].Close()
// make sure dir exists
os.MkdirAll(c.CopyOutDir, 0777)
newDir, err := os.Open(c.CopyOutDir)
if err != nil {
writeErrorC(errC, err)
return
}
defer newDir.Close()
names, err := dir[0].Readdirnames(-1)
if err != nil {
writeErrorC(errC, err)
return
}
for _, n := range names {
if err := copyFileDir(int(dir[0].Fd()), int(newDir.Fd()), n); err != nil {
writeErrorC(errC, err)
return
}
// Link at do not cross device copy, sad..
// if err := unix.Linkat(int(dir[0].Fd()), n, int(newDir.Fd()), n, 0); err != nil {
// writeErrorC(errC, err)
// return
// }
}
}()
}
// wait to finish
wg.Wait()
@ -87,3 +133,29 @@ func copyOutAndCollect(m container.Environment, c *Cmd, ptc []pipeCollector) (ma
return rt, nil
}
func copyFileDir(srcDirFd, dstDirFd int, name string) error {
// open the source file
fd, err := syscall.Openat(srcDirFd, name, syscall.O_RDONLY, 0777)
if err != nil {
return err
}
defer syscall.Close(fd)
var st syscall.Stat_t
if err := syscall.Fstat(fd, &st); err != nil {
return err
}
// open the dst file
dstFd, err := syscall.Openat(dstDirFd, name, syscall.O_WRONLY|syscall.O_CREAT, 0777)
if err != nil {
return err
}
defer syscall.Close(dstFd)
// send file
if _, err := syscall.Sendfile(dstFd, fd, nil, int(st.Size)); err != nil {
return err
}
return nil
}