Go1.9基于http备份文件中心服务器

来源:互联网 发布:聚划算淘宝商城特价区 编辑:程序博客网 时间:2024/05/29 10:24
package mainimport (    "archive/zip"    "crypto/md5"    "encoding/hex"    "encoding/json"    "flag"    "io"    "io/ioutil"    "log"    "net"    "net/http"    "net/url"    "os"    "path/filepath"    "runtime"    "strconv"    "strings"    "time")const (    CodeSuccess = 200    ECodeCreate = 600    ECodeAPPEND = 601    ECodeSave   = 602    ECodeMd5    = 603)var (    cfgPath        string    daemon         bool    sconfig        ServerConfig    cconfig        ClientConfig    uploadFilePath string)func main() {    flag.StringVar(&cfgPath, "c", "ccfg.json", "-c cfg.json 指定配置文件路径")    flag.BoolVar(&daemon, "d", false, "-d 是否以服务端运行")    flag.StringVar(&uploadFilePath, "u", "mount", "-u ./20170803.zip 指定要上传的文件路径")    flag.Parse()    data, err := ioutil.ReadFile(cfgPath)    if err != nil {        log.Fatalf("Read config data error:%s\n", err.Error())    }    if daemon {        err = json.Unmarshal(data, &sconfig)    } else {        err = json.Unmarshal(data, &cconfig)    }    if err != nil {        log.Fatalf("Unmarshal config error:%s\n", err.Error())    }    if daemon {        server()    } else {        client()    }}type ClientConfig struct {    Server        string `json:"server"`    Https         bool   `json:"https"`    Compress      bool   `json:"compress"`    RmCompress    bool   `json:"rmcompress"`    VerifyMd5     bool   `json:"verifymd5"`    CreateDir     bool   `json:"createdir"`    Retry         int    `json:"retry"`    RetryInterval int    `json:"retryinterval"`    User          string `json:"user"`    Password      string `json:"password"`    GameId        string `json:"gameid"`}func client() {    if uploadFilePath == "" {        log.Fatalln("Must specify upload file path")    }    if cconfig.User == "" || cconfig.Password == "" {        log.Fatalln("Must specify authentication user and password")    }    if cconfig.GameId == "" {        log.Fatalf("Must specify gameid")    }    if cconfig.Https {        cconfig.Server = "https://" + cconfig.Server    } else {        cconfig.Server = "http://" + cconfig.Server    }    info, err := os.Lstat(uploadFilePath)    if err != nil {        log.Fatalf("Open file error:%s\n", err.Error())    }    if info.IsDir() {        log.Fatalf("Upload path:%s is dirctory\n", uploadFilePath)    }    if cconfig.Compress {        var err error        uploadFilePath, err = compress(uploadFilePath)        if err != nil {            log.Fatalf("Compress file error:%s\n", err.Error())        }        if cconfig.RmCompress {            defer os.Remove(uploadFilePath)        }    }    var md5str string    if cconfig.VerifyMd5 {        md5str = md5sum(uploadFilePath)        if md5str == "" {            log.Fatalf("Get %s md5 error:%s\n", uploadFilePath, md5str)        } else {            log.Printf("%s md5:%s\n", uploadFilePath, md5str)        }    }    File, err := os.Open(uploadFilePath)    if err != nil {        log.Fatalf("Open file error:%s\n", err.Error())    }    //  defer File.Close()  http请求结束后会调用Close()方法    info, err = File.Stat()    if err != nil {        log.Fatalf("Get file info error:%s\n", err.Error())    }    req, err := http.NewRequest("POST", cconfig.Server+"/upload", File)    if err != nil {        log.Fatalf("Init request error:%s\n", err.Error())    }    req.SetBasicAuth(cconfig.User, cconfig.Password)    req.ContentLength = info.Size()    if cconfig.CreateDir {        req.Header.Set("path", uploadFilePath)    } else {        req.Header.Set("path", filepath.Base(uploadFilePath))    }    req.Header.Set("id", cconfig.GameId)    req.Header.Set("md5", md5str)    doReq(req, uploadFilePath)}func doReq(req *http.Request, uploadFilePath string) {    var exit bool = false    var filesize = req.ContentLengthreupload:    resp, err := http.DefaultClient.Do(req)    if err != nil {        log.Printf("Upload file error:%s\n", err.Error())        ue, ok := err.(*url.Error)        if !ok {            return        }    retry:        if !ue.Temporary() || cconfig.Retry <= 0 {            if cconfig.Retry <= 0 {                return            }            if _, ok = ue.Err.(*net.OpError); !ok {                if !strings.Contains(ue.Err.Error(), "An existing connection was forcibly closed by the remote host") {                    return                }            }        }        cconfig.Retry--        time.Sleep(time.Second * time.Duration(cconfig.RetryInterval))        req.Body = nil        req.ContentLength = 0        req.Header.Set("retry", "true")        resp, err := http.DefaultClient.Do(req)        if err != nil {            goto retry        }        var (            size    int64  = 0            sizestr string = resp.Header.Get("size")        )        if sizestr != "0" {            size, err = strconv.ParseInt(sizestr, 10, 0)            if err != nil {                sizestr = "0"            }        }        log.Printf("File:%s,Already send %s\n", uploadFilePath, sizestr)        File, err := os.Open(uploadFilePath)        if err != nil {            log.Printf("Open file error:%s\n", err.Error())            return        }        File.Seek(size, 0)        req.Header.Del("retry")        req.Header.Set("size", sizestr)        req.ContentLength = filesize - size        req.Body = File        goto reupload    }    switch resp.StatusCode {    case CodeSuccess:        log.Printf("File %s upload successful\n", uploadFilePath)    case ECodeAPPEND, ECodeSave:        log.Printf("Retry upload %s error\n", uploadFilePath)        if !exit {            File, err := os.Open(uploadFilePath)            if err != nil {                log.Printf("Open file error:%s\n", err.Error())                return            }            File.Seek(0, 0)            req.Header.Del("retry")            req.Header.Set("size", "")            req.ContentLength = filesize            req.Body = File            exit = true            goto reupload        }    case ECodeCreate:        log.Printf("Upload %s error\n", uploadFilePath)    case ECodeMd5:        log.Printf("Upload %s md5sum not match\n", uploadFilePath)    default:        log.Printf("Is undefind statuscode %d\n", resp.StatusCode)    }}type ServerConfig struct {    Listen   string                       `json:"listen"`    RootPath string                       `json:"rootpath"`    User     map[string]map[string]string `json:"user"`    Key      string                       `json:"key"`    Crt      string                       `json:"crt"`}func server() {    if sconfig.RootPath == "" {        var err error        if sconfig.RootPath, err = os.Getwd(); err != nil {            log.Fatalf("Get current dirpath error:%s\n", err.Error())        }    }    stat, err := os.Lstat(sconfig.RootPath)    if err != nil {        log.Fatalf("List RootPath stat error:%s\n", err.Error())    }    if !stat.IsDir() {        log.Fatalf("RootPath:%s must directory path\n", sconfig.RootPath)    }    if runtime.GOOS == "windows" {        sconfig.RootPath = strings.Replace(sconfig.RootPath, "\\", "/", -1)    }    http.HandleFunc("/", route)    if sconfig.Crt != "" && sconfig.Key != "" {        err = http.ListenAndServeTLS(sconfig.Listen, sconfig.Crt, sconfig.Key, nil)    } else {        err = http.ListenAndServe(sconfig.Listen, nil)    }    if err != nil {        log.Printf("listen %s error:%s\n", sconfig.Listen, err.Error())    }}func route(w http.ResponseWriter, r *http.Request) {    log.Printf("RemoteAddr:%s URI:%s\n", r.RemoteAddr, r.RequestURI)    defer r.Body.Close()    var code int = http.StatusOK    switch r.URL.Path {    case "/upload":        rootpath, size := baseInfo(r)        if rootpath == "" {            http.Error(w, "Authentication baseinfo failed", http.StatusForbidden)            return        }        log.Printf("FilePath:%s Size:%d\n", rootpath, size)        code = upload(w, r, rootpath, size)    case "/config":        if strings.Index(r.RemoteAddr, "127.0.0.1") == 0 {            buf, _ := json.Marshal(sconfig)            w.WriteHeader(code)            w.Write(buf)        } else {            w.WriteHeader(http.StatusForbidden)        }        return    case "/flush":        if strings.Index(r.RemoteAddr, "127.0.0.1") == 0 {            buf, err := ioutil.ReadFile(cfgPath)            if err == nil {                var cfg ServerConfig                err = json.Unmarshal(buf, &cfg)                if err == nil {                    sconfig = cfg                    w.WriteHeader(CodeSuccess)                    w.Write(buf)                    return                } else {                    code = http.StatusInternalServerError                    log.Printf("Read config error:%s\n", err.Error())                }            } else {                code = http.StatusInternalServerError                log.Printf("Unmarshal config error:" + err.Error())            }        } else {            code = http.StatusForbidden        }    default:        code = http.StatusNotFound    }    w.WriteHeader(code)}func upload(w http.ResponseWriter, r *http.Request, rootpath string, size int64) int {    var (        err  error        code int        File *os.File    )    if r.Header.Get("retry") == "true" {        stat, err := os.Lstat(rootpath + ".tmp")        if err == nil {            w.Header().Set("size", strconv.FormatInt(stat.Size(), 10))        } else {            w.Header().Set("size", "0")        }        return CodeSuccess    }    os.MkdirAll(filepath.Dir(rootpath), 0644)    //如果size不为0则认为是续传    if size == 0 {        File, err = os.OpenFile(rootpath+".tmp", os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0644)        code = ECodeCreate    } else {        File, err = os.OpenFile(rootpath+".tmp", os.O_APPEND|os.O_RDWR, 0644)        code = ECodeAPPEND    }    if err != nil {        log.Printf("Create file error:%s\n", err.Error())        return code    }    info, err := File.Stat()    if err != nil {        log.Printf("Get Fileinfo error:%s\n", err.Error())        return ECodeCreate    }    //查看文件size是否和请求size一致,如果不一致则返回错误    if info.Size() != size {        return ECodeAPPEND    }    _, err = io.Copy(File, r.Body)    File.Close()    if err != nil {        log.Printf("Save body faild:%s\n", err.Error())        return ECodeSave    }    var m, path string    md5str := strings.ToLower(r.Header.Get("md5"))    if md5str != "" {        if m = md5sum(rootpath + ".tmp"); md5str != m {            return ECodeMd5        }    } else {        m = time.Now().Format("20060102151605")    }    index := strings.LastIndex(rootpath, ".")    if index != -1 {        path = string(rootpath[:index] + "_" + m + string(rootpath[index:]))    } else {        path = rootpath    }    os.Rename(rootpath+".tmp", path)    return CodeSuccess}func baseInfo(r *http.Request) (string, int64) {    id := basicAuth(r)    if id == "" {        return "", 0    }    var path = r.Header.Get("path")    if path = scopePath(path); path == "" {        return "", 0    }    DirPath := sconfig.RootPath + "/" + id + "/" + time.Now().Format("20060102") + "/"    info, err := os.Lstat(DirPath)    if err != nil {        if os.IsNotExist(err) {            os.MkdirAll(DirPath, 0644)        } else {            log.Printf("Check dirpath error:%s\n", err.Error())            return "", 0        }    } else if !info.IsDir() {        return "", 0    }    path = DirPath + path    if size := r.Header.Get("size"); size != "" && size != "0" {        Size, err := strconv.ParseInt(size, 10, 0)        if err != nil {            log.Printf("Parse Size error:%s\n", err.Error())            return "", 0        }        return path, Size    }    return path, 0}func basicAuth(r *http.Request) string {    u, p, ok := r.BasicAuth()    if !ok {        log.Println("Get basic auth error.")        return ""    }    id := r.Header.Get("id")    if sconfig.User[id][u] != p {        log.Printf("user and password unmatch:%s %s\n", u, p)        return ""    }    return id}func scopePath(filePath string) string {    if filePath == "" {        return filePath    }    filePath = strings.Replace(filePath, "\\", "/", -1)    switch runtime.GOOS {    case "linux":        if filePath[0] == '/' {            filePath = filePath[1:]        }    case "windows":        if l := strings.Index(filePath, ":"); l != -1 {            filePath = filePath[l+1:]        }        filePath = strings.TrimLeft(filePath, "/")    }    return strings.TrimLeft(filePath, "./")}func md5sum(filePath string) string {    File, err := os.Open(filePath)    if err != nil {        return ""    }    defer File.Close()    sum := md5.New()    io.Copy(sum, File)    m := make([]byte, 0, 32)    return hex.EncodeToString(sum.Sum(m))}const zone = 8 * 60 * 60func compress(path string) (string, error) {    index := strings.LastIndex(path, ".")    if index < 0 {        index = len(path)    }    fpath := string(path[:index]) + ".zip"    if fpath == path {        return path, nil    }    sFile, err := os.Open(path)    if err != nil {        return "", err    }    defer sFile.Close()    info, err := sFile.Stat()    if err != nil {        return "", err    }    File, err := os.Create(fpath)    if err != nil {        return "", err    }    defer File.Close()    w := zip.NewWriter(File)    defer w.Close()    header, err := zip.FileInfoHeader(info)    if err != nil {        return "", err    }    header.Method = zip.Deflate    header.SetModTime(time.Unix(info.ModTime().Unix()+zone, 0))    f, err := w.CreateHeader(header)    if err != nil {        return "", err    }    _, err = io.Copy(f, sFile)    if err != nil {        return "", err    }    return fpath, nil}