github.com/annwntech/go-micro/v2@v2.9.5/util/file/handler.go (about)

     1  package file
     2  
     3  import (
     4  	"io"
     5  	"os"
     6  	"path/filepath"
     7  	"sync"
     8  
     9  	"github.com/annwntech/go-micro/v2/errors"
    10  	"github.com/annwntech/go-micro/v2/logger"
    11  	"github.com/annwntech/go-micro/v2/server"
    12  	proto "github.com/annwntech/go-micro/v2/util/file/proto"
    13  	"golang.org/x/net/context"
    14  )
    15  
    16  // NewHandler is a handler that can be registered with a micro Server
    17  func NewHandler(readDir string) proto.FileHandler {
    18  	return &handler{
    19  		readDir: readDir,
    20  		session: &session{
    21  			files: make(map[int64]*os.File),
    22  		},
    23  	}
    24  }
    25  
    26  // RegisterHandler is a convenience method for registering a handler
    27  func RegisterHandler(s server.Server, readDir string) {
    28  	proto.RegisterFileHandler(s, NewHandler(readDir))
    29  }
    30  
    31  type handler struct {
    32  	readDir string
    33  	session *session
    34  }
    35  
    36  func (h *handler) Open(ctx context.Context, req *proto.OpenRequest, rsp *proto.OpenResponse) error {
    37  	path := filepath.Join(h.readDir, req.Filename)
    38  	flags := os.O_CREATE | os.O_RDWR
    39  	if req.GetTruncate() {
    40  		flags = flags | os.O_TRUNC
    41  	}
    42  	file, err := os.OpenFile(path, flags, 0666)
    43  	if err != nil {
    44  		return errors.InternalServerError("go.micro.server", err.Error())
    45  	}
    46  
    47  	rsp.Id = h.session.Add(file)
    48  	rsp.Result = true
    49  
    50  	logger.Debugf("Open %s, sessionId=%d", req.Filename, rsp.Id)
    51  
    52  	return nil
    53  }
    54  
    55  func (h *handler) Close(ctx context.Context, req *proto.CloseRequest, rsp *proto.CloseResponse) error {
    56  	h.session.Delete(req.Id)
    57  	logger.Debugf("Close sessionId=%d", req.Id)
    58  	return nil
    59  }
    60  
    61  func (h *handler) Stat(ctx context.Context, req *proto.StatRequest, rsp *proto.StatResponse) error {
    62  	path := filepath.Join(h.readDir, req.Filename)
    63  	fi, err := os.Stat(path)
    64  	if os.IsNotExist(err) {
    65  		return errors.InternalServerError("go.micro.srv.file", err.Error())
    66  	}
    67  
    68  	if fi.IsDir() {
    69  		rsp.Type = "Directory"
    70  	} else {
    71  		rsp.Type = "File"
    72  		rsp.Size = fi.Size()
    73  	}
    74  
    75  	rsp.LastModified = fi.ModTime().Unix()
    76  	logger.Debugf("Stat %s, %#v", req.Filename, rsp)
    77  
    78  	return nil
    79  }
    80  
    81  func (h *handler) Read(ctx context.Context, req *proto.ReadRequest, rsp *proto.ReadResponse) error {
    82  	file := h.session.Get(req.Id)
    83  	if file == nil {
    84  		return errors.InternalServerError("go.micro.srv.file", "You must call open first.")
    85  	}
    86  
    87  	rsp.Data = make([]byte, req.Size)
    88  	n, err := file.ReadAt(rsp.Data, req.Offset)
    89  	if err != nil && err != io.EOF {
    90  		return errors.InternalServerError("go.micro.srv.file", err.Error())
    91  	}
    92  
    93  	if err == io.EOF {
    94  		rsp.Eof = true
    95  	}
    96  
    97  	rsp.Size = int64(n)
    98  	rsp.Data = rsp.Data[:n]
    99  
   100  	logger.Debugf("Read sessionId=%d, Offset=%d, n=%d", req.Id, req.Offset, rsp.Size)
   101  
   102  	return nil
   103  }
   104  
   105  func (h *handler) Write(ctx context.Context, req *proto.WriteRequest, rsp *proto.WriteResponse) error {
   106  	file := h.session.Get(req.Id)
   107  	if file == nil {
   108  		return errors.InternalServerError("go.micro.srv.file", "You must call open first.")
   109  	}
   110  
   111  	if _, err := file.WriteAt(req.GetData(), req.GetOffset()); err != nil {
   112  		return err
   113  	}
   114  
   115  	logger.Debugf("Write sessionId=%d, Offset=%d, n=%d", req.Id, req.Offset)
   116  
   117  	return nil
   118  }
   119  
   120  type session struct {
   121  	sync.Mutex
   122  	files   map[int64]*os.File
   123  	counter int64
   124  }
   125  
   126  func (s *session) Add(file *os.File) int64 {
   127  	s.Lock()
   128  	defer s.Unlock()
   129  
   130  	s.counter += 1
   131  	s.files[s.counter] = file
   132  
   133  	return s.counter
   134  }
   135  
   136  func (s *session) Get(id int64) *os.File {
   137  	s.Lock()
   138  	defer s.Unlock()
   139  	return s.files[id]
   140  }
   141  
   142  func (s *session) Delete(id int64) {
   143  	s.Lock()
   144  	defer s.Unlock()
   145  
   146  	if file, exist := s.files[id]; exist {
   147  		file.Close()
   148  		delete(s.files, id)
   149  	}
   150  }
   151  
   152  func (s *session) Len() int {
   153  	return len(s.files)
   154  }