handler.go
1 package fileserver 2 3 import ( 4 "encoding/base64" 5 "encoding/json" 6 "io" 7 "log/slog" 8 "os" 9 "path/filepath" 10 "sort" 11 "strings" 12 "sync" 13 ) 14 15 var protectedPaths = map[string]bool{ 16 "/": true, "/bin": true, "/sbin": true, "/etc": true, 17 "/usr": true, "/var": true, "/lib": true, "/lib64": true, 18 "/boot": true, "/proc": true, "/sys": true, "/dev": true, 19 "/home": true, "/root": true, "/tmp": true, "/run": true, 20 "/opt": true, "/srv": true, "/mnt": true, "/media": true, 21 } 22 23 type uploadState struct { 24 file *os.File 25 tempPath string 26 destPath string 27 received int64 28 expected int64 29 } 30 31 type Handler struct { 32 config *SecurityConfig 33 uploads map[string]*uploadState 34 uploadsMu sync.Mutex 35 } 36 37 const ChunkSize = 256 * 1024 38 39 func NewHandler(config *SecurityConfig) *Handler { 40 return &Handler{ 41 config: config, 42 uploads: make(map[string]*uploadState), 43 } 44 } 45 46 func (h *Handler) HandleMessage(data []byte, send func(interface{})) { 47 var raw map[string]interface{} 48 if err := json.Unmarshal(data, &raw); err != nil { 49 slog.Error("failed to decode file message", "error", err) 50 return 51 } 52 53 msgType, _ := raw["type"].(string) 54 reqID, _ := raw["requestId"].(string) 55 56 if !h.config.Enabled { 57 send(FileOperationResult{ 58 Type: "file_result", 59 RequestID: reqID, 60 Success: false, 61 Error: "file server is disabled", 62 }) 63 return 64 } 65 66 switch msgType { 67 case "file_list": 68 h.handleList(data, reqID, send) 69 case "file_download": 70 h.handleDownload(data, reqID, send) 71 case "file_upload_start": 72 h.handleUploadStart(data, reqID, send) 73 case "file_upload_chunk": 74 h.handleUploadChunk(data, reqID, send) 75 case "file_mkdir": 76 h.handleMkdir(data, reqID, send) 77 case "file_delete": 78 h.handleDelete(data, reqID, send) 79 case "file_rename": 80 h.handleRename(data, reqID, send) 81 default: 82 send(FileOperationResult{ 83 Type: "file_result", 84 RequestID: reqID, 85 Success: false, 86 Error: "unknown file operation type", 87 }) 88 } 89 } 90 91 func (h *Handler) handleList(data []byte, reqID string, send func(interface{})) { 92 var req FileListRequest 93 if err := json.Unmarshal(data, &req); err != nil { 94 h.sendError(reqID, "invalid list request", send) 95 return 96 } 97 98 absPath, err := h.config.ValidatePath(req.Path) 99 if err != nil { 100 h.sendError(reqID, err.Error(), send) 101 return 102 } 103 104 entries, err := os.ReadDir(absPath) 105 if err != nil { 106 h.sendError(reqID, err.Error(), send) 107 return 108 } 109 110 var fileInfos []FileInfo 111 for _, entry := range entries { 112 if len(fileInfos) >= h.config.MaxListEntries { 113 break 114 } 115 info, err := entry.Info() 116 if err != nil { 117 continue 118 } 119 fileInfos = append(fileInfos, FileInfo{ 120 Name: info.Name(), 121 IsDir: info.IsDir(), 122 Size: info.Size(), 123 Mode: info.Mode().String(), 124 ModTime: info.ModTime().UnixMilli(), 125 }) 126 } 127 128 sort.Slice(fileInfos, func(i, j int) bool { 129 if fileInfos[i].IsDir != fileInfos[j].IsDir { 130 return fileInfos[i].IsDir 131 } 132 return strings.ToLower(fileInfos[i].Name) < strings.ToLower(fileInfos[j].Name) 133 }) 134 135 send(FileListResponse{ 136 Type: "file_list_result", 137 RequestID: reqID, 138 Success: true, 139 Path: req.Path, 140 Entries: fileInfos, 141 }) 142 } 143 144 func (h *Handler) handleDownload(data []byte, reqID string, send func(interface{})) { 145 var req FileDownloadRequest 146 if err := json.Unmarshal(data, &req); err != nil { 147 h.sendError(reqID, "invalid download request", send) 148 return 149 } 150 151 absPath, err := h.config.ValidatePath(req.Path) 152 if err != nil { 153 h.sendError(reqID, err.Error(), send) 154 return 155 } 156 157 info, err := os.Stat(absPath) 158 if err != nil { 159 h.sendError(reqID, err.Error(), send) 160 return 161 } 162 if info.IsDir() { 163 h.sendError(reqID, "cannot download directory", send) 164 return 165 } 166 if err := h.config.ValidateFileSize(info.Size()); err != nil { 167 h.sendError(reqID, err.Error(), send) 168 return 169 } 170 171 file, err := os.Open(absPath) 172 if err != nil { 173 h.sendError(reqID, err.Error(), send) 174 return 175 } 176 defer file.Close() 177 178 buf := make([]byte, ChunkSize) 179 index := 0 180 totalSize := info.Size() 181 182 for { 183 n, err := file.Read(buf) 184 if n > 0 { 185 // Check if this is the last chunk by seeing if we've read to EOF 186 // or if the next read would be EOF 187 isFinal := err == io.EOF 188 if !isFinal && n < ChunkSize { 189 // If we read less than a full chunk and no error, peek ahead 190 // This handles cases where Read returns partial data before EOF 191 isFinal = true 192 } 193 chunk := FileDownloadChunk{ 194 Type: "file_download_chunk", 195 RequestID: reqID, 196 Index: index, 197 Data: base64.StdEncoding.EncodeToString(buf[:n]), 198 Final: isFinal, 199 } 200 if index == 0 { 201 chunk.TotalSize = totalSize 202 } 203 send(chunk) 204 index++ 205 206 if isFinal { 207 break 208 } 209 } 210 if err != nil { 211 if err != io.EOF { 212 h.sendError(reqID, "read error: "+err.Error(), send) 213 } else if index == 0 { 214 // Empty file case 215 send(FileDownloadChunk{ 216 Type: "file_download_chunk", 217 RequestID: reqID, 218 Index: 0, 219 Data: "", 220 Final: true, 221 TotalSize: 0, 222 }) 223 } 224 break 225 } 226 } 227 } 228 229 func (h *Handler) handleUploadStart(data []byte, reqID string, send func(interface{})) { 230 var req FileUploadStartRequest 231 if err := json.Unmarshal(data, &req); err != nil { 232 h.sendError(reqID, "invalid upload start request", send) 233 return 234 } 235 236 absPath, err := h.config.ValidateWrite(req.Path) 237 if err != nil { 238 h.sendError(reqID, err.Error(), send) 239 return 240 } 241 if err := h.config.ValidateFileSize(req.Size); err != nil { 242 h.sendError(reqID, err.Error(), send) 243 return 244 } 245 246 tempFile, err := os.CreateTemp(filepath.Dir(absPath), ".easyshell_upload_"+reqID+"_*") 247 if err != nil { 248 h.sendError(reqID, "failed to create temp file: "+err.Error(), send) 249 return 250 } 251 252 h.uploadsMu.Lock() 253 h.uploads[reqID] = &uploadState{ 254 file: tempFile, 255 tempPath: tempFile.Name(), 256 destPath: absPath, 257 expected: req.Size, 258 } 259 h.uploadsMu.Unlock() 260 261 send(FileOperationResult{Type: "file_result", RequestID: reqID, Success: true}) 262 } 263 264 func (h *Handler) handleUploadChunk(data []byte, reqID string, send func(interface{})) { 265 var req FileUploadChunkRequest 266 if err := json.Unmarshal(data, &req); err != nil { 267 h.sendError(reqID, "invalid upload chunk request", send) 268 return 269 } 270 271 h.uploadsMu.Lock() 272 state, ok := h.uploads[reqID] 273 h.uploadsMu.Unlock() 274 275 if !ok { 276 h.sendError(reqID, "upload session not found", send) 277 return 278 } 279 280 decoded, err := base64.StdEncoding.DecodeString(req.Data) 281 if err != nil { 282 h.sendError(reqID, "invalid base64 chunk", send) 283 return 284 } 285 286 n, err := state.file.Write(decoded) 287 if err != nil { 288 h.sendError(reqID, "write error: "+err.Error(), send) 289 return 290 } 291 state.received += int64(n) 292 293 if req.Final { 294 state.file.Close() 295 h.uploadsMu.Lock() 296 delete(h.uploads, reqID) 297 h.uploadsMu.Unlock() 298 299 if err := os.Rename(state.tempPath, state.destPath); err != nil { 300 h.sendError(reqID, "failed to commit file: "+err.Error(), send) 301 return 302 } 303 } 304 305 send(FileOperationResult{Type: "file_result", RequestID: reqID, Success: true}) 306 } 307 308 func (h *Handler) handleMkdir(data []byte, reqID string, send func(interface{})) { 309 var req FileMkdirRequest 310 if err := json.Unmarshal(data, &req); err != nil { 311 h.sendError(reqID, "invalid mkdir request", send) 312 return 313 } 314 315 absPath, err := h.config.ValidateWrite(req.Path) 316 if err != nil { 317 h.sendError(reqID, err.Error(), send) 318 return 319 } 320 321 if err := os.MkdirAll(absPath, 0755); err != nil { 322 h.sendError(reqID, err.Error(), send) 323 return 324 } 325 326 send(FileOperationResult{Type: "file_result", RequestID: reqID, Success: true}) 327 } 328 329 func (h *Handler) handleDelete(data []byte, reqID string, send func(interface{})) { 330 var req FileDeleteRequest 331 if err := json.Unmarshal(data, &req); err != nil { 332 h.sendError(reqID, "invalid delete request", send) 333 return 334 } 335 336 absPath, err := h.config.ValidateWrite(req.Path) 337 if err != nil { 338 h.sendError(reqID, err.Error(), send) 339 return 340 } 341 342 if protectedPaths[absPath] { 343 h.sendError(reqID, "cannot delete system protected directory", send) 344 return 345 } 346 347 if err := os.RemoveAll(absPath); err != nil { 348 h.sendError(reqID, err.Error(), send) 349 return 350 } 351 352 send(FileOperationResult{Type: "file_result", RequestID: reqID, Success: true}) 353 } 354 355 func (h *Handler) handleRename(data []byte, reqID string, send func(interface{})) { 356 var req FileRenameRequest 357 if err := json.Unmarshal(data, &req); err != nil { 358 h.sendError(reqID, "invalid rename request", send) 359 return 360 } 361 362 oldAbs, err := h.config.ValidateWrite(req.OldPath) 363 if err != nil { 364 h.sendError(reqID, "old path: "+err.Error(), send) 365 return 366 } 367 newAbs, err := h.config.ValidateWrite(req.NewPath) 368 if err != nil { 369 h.sendError(reqID, "new path: "+err.Error(), send) 370 return 371 } 372 373 if protectedPaths[oldAbs] { 374 h.sendError(reqID, "cannot rename system protected directory", send) 375 return 376 } 377 378 if err := os.Rename(oldAbs, newAbs); err != nil { 379 h.sendError(reqID, err.Error(), send) 380 return 381 } 382 383 send(FileOperationResult{Type: "file_result", RequestID: reqID, Success: true}) 384 } 385 386 func (h *Handler) sendError(reqID string, errStr string, send func(interface{})) { 387 send(FileOperationResult{ 388 Type: "file_result", 389 RequestID: reqID, 390 Success: false, 391 Error: errStr, 392 }) 393 } 394 395 func (h *Handler) CleanupStaleTempFiles() { 396 if !h.config.Enabled { 397 return 398 } 399 root := h.config.RootPath 400 _ = filepath.Walk(root, func(path string, info os.FileInfo, err error) error { 401 if err != nil { 402 return filepath.SkipDir 403 } 404 if info.IsDir() { 405 if path != root { 406 return filepath.SkipDir 407 } 408 return nil 409 } 410 if strings.HasPrefix(filepath.Base(path), ".easyshell_upload_") { 411 slog.Info("cleaning up stale temp file", "path", path) 412 _ = os.Remove(path) 413 } 414 return nil 415 }) 416 }