Skip to content

Commit

Permalink
add unexpected eof handler
Browse files Browse the repository at this point in the history
  • Loading branch information
Son Roy Almerol committed Nov 12, 2024
1 parent 0e13141 commit 7b68f5e
Show file tree
Hide file tree
Showing 5 changed files with 40 additions and 16 deletions.
3 changes: 0 additions & 3 deletions internal/agent/sftp/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,9 +37,6 @@ func (s *SFTPConfig) GetRegistryKey() string {

func InitializeSFTPConfig(svc service.Service, driveLetter string) (*SFTPConfig, error) {
var err error
if err != nil {
return nil, fmt.Errorf("InitializeLogger: failed to initialize logger -> %w", err)
}

baseKey, _, err := registry.CreateKey(registry.LOCAL_MACHINE, "Software\\PBSPlus\\Config", registry.QUERY_VALUE)
if err != nil {
Expand Down
28 changes: 28 additions & 0 deletions internal/agent/sftp/eof.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
package sftp

import "io"

type eofInjectingReaderAt struct {
file io.ReaderAt
length int64
}

func NewEOFInjectingReaderAt(file io.ReaderAt, length int64) io.ReaderAt {
return &eofInjectingReaderAt{file: file, length: length}
}

func (r *eofInjectingReaderAt) ReadAt(p []byte, off int64) (n int, err error) {
if off >= r.length {
// If reading beyond the file length, return EOF to simulate a proper end.
return 0, io.EOF
}

// Attempt to read as usual.
n, err = r.file.ReadAt(p, off)
if err == io.ErrUnexpectedEOF || n < len(p) {
// If an unexpected EOF or partial read, return what we have and inject EOF.
return n, io.EOF
}

return n, err
}
12 changes: 2 additions & 10 deletions internal/agent/sftp/filelister.go
Original file line number Diff line number Diff line change
Expand Up @@ -88,8 +88,8 @@ func (h *SftpHandler) setFilePath(r *sftp.Request) {
r.Filepath = filepath.Join(h.Snapshot.SnapshotPath, filepath.Clean(r.Filepath))
}

func (h *SftpHandler) fetch(path string, mode int) (*os.File, error) {
return os.OpenFile(path, mode, 0777)
func (h *SftpHandler) fetch(path string) (*os.File, error) {
return os.Open(path)
}

func wildcardToRegex(pattern string) string {
Expand Down Expand Up @@ -175,13 +175,5 @@ func skipFile(path string, fileInfo os.FileInfo) bool {
}
}

if !fileInfo.IsDir() {
f, err := os.Open(path)
if err != nil {
return true
}
f.Close()
}

return false
}
11 changes: 9 additions & 2 deletions internal/agent/sftp/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -40,12 +40,19 @@ func (h *SftpHandler) Fileread(r *sftp.Request) (io.ReaderAt, error) {
h.Snapshot.UpdateTimestamp()
h.setFilePath(r)

file, err := h.fetch(r.Filepath, os.O_RDONLY)
file, err := h.fetch(r.Filepath)
if err != nil {
log.Printf("error reading file: %v", err)
return nil, err
}
return file, nil

stat, err := os.Lstat(r.Filepath)
if err != nil {
log.Printf("error getting file stats: %v", err)
return nil, err
}

return NewEOFInjectingReaderAt(file, stat.Size()), nil
}

func (h *SftpHandler) Filewrite(r *sftp.Request) (io.WriterAt, error) {
Expand Down
2 changes: 1 addition & 1 deletion internal/backend/mount/mount.go
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ func Mount(target *store.Target) (*AgentMount, error) {
"--sftp-user", "proxmox",
"--sftp-host", agentHost,
"--allow-other",
"--sftp-shell-type", "unix",
"--sftp-shell-type", "none",
":sftp:/", agentMount.Path,
}

Expand Down

0 comments on commit 7b68f5e

Please sign in to comment.