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

package main

import (
    "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.ContentLength
reupload:
    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 * 60

func 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
}