diff options
author | Quentin Dufour <quentin@deuxfleurs.fr> | 2021-11-19 19:54:49 +0100 |
---|---|---|
committer | Quentin Dufour <quentin@deuxfleurs.fr> | 2021-11-19 19:54:49 +0100 |
commit | 0ee29e31ddcc81f541de7459b0a5e40dfa552672 (patch) | |
tree | 859ff133f8c78bd034b0c2184cdad0ce9f38b065 /sftp/allocator.go | |
parent | 93631b4e3d5195d446504db1c4a2bc7468b3ef28 (diff) | |
download | bagage-0ee29e31ddcc81f541de7459b0a5e40dfa552672.tar.gz bagage-0ee29e31ddcc81f541de7459b0a5e40dfa552672.zip |
Working on SFTP
Diffstat (limited to 'sftp/allocator.go')
-rw-r--r-- | sftp/allocator.go | 101 |
1 files changed, 101 insertions, 0 deletions
diff --git a/sftp/allocator.go b/sftp/allocator.go new file mode 100644 index 0000000..fc1b6f0 --- /dev/null +++ b/sftp/allocator.go @@ -0,0 +1,101 @@ +package sftp + +/* + Imported from: https://github.com/pkg/sftp + */ + +import ( + "sync" +) + +type allocator struct { + sync.Mutex + available [][]byte + // map key is the request order + used map[uint32][][]byte +} + +func newAllocator() *allocator { + return &allocator{ + // micro optimization: initialize available pages with an initial capacity + available: make([][]byte, 0, SftpServerWorkerCount*2), + used: make(map[uint32][][]byte), + } +} + +// GetPage returns a previously allocated and unused []byte or create a new one. +// The slice have a fixed size = maxMsgLength, this value is suitable for both +// receiving new packets and reading the files to serve +func (a *allocator) GetPage(requestOrderID uint32) []byte { + a.Lock() + defer a.Unlock() + + var result []byte + + // get an available page and remove it from the available ones. + if len(a.available) > 0 { + truncLength := len(a.available) - 1 + result = a.available[truncLength] + + a.available[truncLength] = nil // clear out the internal pointer + a.available = a.available[:truncLength] // truncate the slice + } + + // no preallocated slice found, just allocate a new one + if result == nil { + result = make([]byte, maxMsgLength) + } + + // put result in used pages + a.used[requestOrderID] = append(a.used[requestOrderID], result) + + return result +} + +// ReleasePages marks unused all pages in use for the given requestID +func (a *allocator) ReleasePages(requestOrderID uint32) { + a.Lock() + defer a.Unlock() + + if used := a.used[requestOrderID]; len(used) > 0 { + a.available = append(a.available, used...) + } + delete(a.used, requestOrderID) +} + +// Free removes all the used and available pages. +// Call this method when the allocator is not needed anymore +func (a *allocator) Free() { + a.Lock() + defer a.Unlock() + + a.available = nil + a.used = make(map[uint32][][]byte) +} + +func (a *allocator) countUsedPages() int { + a.Lock() + defer a.Unlock() + + num := 0 + for _, p := range a.used { + num += len(p) + } + return num +} + +func (a *allocator) countAvailablePages() int { + a.Lock() + defer a.Unlock() + + return len(a.available) +} + +func (a *allocator) isRequestOrderIDUsed(requestOrderID uint32) bool { + a.Lock() + defer a.Unlock() + + _, ok := a.used[requestOrderID] + return ok +} + |