Golang笔记

golang + bolt 实现的短网址微服务

bolt 是高性能的kv 嵌入式数据库,很适合用在微服务上,下面是goji + bolt 写的短网址服务。

package main

import (
	"crypto/md5"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strconv"

	"github.com/boltdb/bolt"
	"goji.io"
	"goji.io/pat"
)

type DBHandler struct {
	DB *bolt.DB
}

func (ah *DBHandler) gotoshorturl(w http.ResponseWriter, r *http.Request) {
	shorturl := pat.Param(r, "shorturl")

	if len(shorturl) != 6 {
		fmt.Fprint(w, "404: not found a!")
		return
	}

	var longURL string
	ah.DB.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("shortURL"))
		v := b.Get([]byte(shorturl))
		if len(v) > 0 {
			longURL = string(v)
		}
		return nil
	})

	if len(longURL) > 0 {
		http.Redirect(w, r, longURL, http.StatusSeeOther)
		return
	}
	fmt.Fprint(w, "404: not found b!")

}

func (ah *DBHandler) api(w http.ResponseWriter, r *http.Request) {

	msg := struct {
		Code   string
		Result string
		Msg    string
	}{}

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")

	err := r.ParseForm()
	if err != nil {
		msg.Code = "1"
		msg.Msg = err.Error()
		json.NewEncoder(w).Encode(msg)
		return
	}

	longURL := r.FormValue("longurl")

	if len(longURL) == 0 {
		msg.Code = "2"
		msg.Msg = "url len is 0!"
		json.NewEncoder(w).Encode(msg)
		return
	} else {
		_, err := url.ParseRequestURI(longURL)
		if err != nil {
			msg.Code = "3"
			msg.Msg = "url validating fail!"
			json.NewEncoder(w).Encode(msg)
			return
		}
	}

	urlList := URLshorten(longURL)

	curSURL := ""
	surlExist := false // 相同的longURL 已存在
	hasSURL := false   // 还有坑
	surlMap := make(map[string]int)

	if len(urlList) > 0 {
		ah.DB.View(func(tx *bolt.Tx) error {
			b := tx.Bucket([]byte("shortURL"))

			for _, surl := range urlList {
				curSURL = surl
				v := b.Get([]byte(surl))

				vs := string(v)
				if len(v) > 0 {
					if vs == longURL {
						surlMap[surl] = 2
						surlExist = true
						break
					} else {
						surlMap[surl] = 1
					}
				} else {
					surlMap[surl] = 0
					hasSURL = true
				}
			}

			return nil
		})
	}

	if surlExist {
		msg.Code = "0"
		msg.Msg = "ok"
		msg.Result = fmt.Sprintf("http://%s/%s", r.Host, curSURL)
		json.NewEncoder(w).Encode(msg)
		return
	}

	if hasSURL {
		ah.DB.Update(func(tx *bolt.Tx) error {
			b := tx.Bucket([]byte("shortURL"))
			for k, v := range surlMap {
				if v == 0 {
					b.Put([]byte(k), []byte(longURL))
					msg.Code = "0"
					msg.Msg = "ok"
					msg.Result = fmt.Sprintf("http://%s/%s", r.Host, k)
					break
				}
			}
			return nil
		})
	} else {
		msg.Code = "4"
		msg.Msg = "no place"
	}

	json.NewEncoder(w).Encode(msg)

}

func main() {
	port := flag.String("port", "8000", "Port for server")
	dbfile := flag.String("dbfile", "my.db", "full path of db file")
	flag.Parse()

	db, err := bolt.Open(*dbfile, 0600, nil)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()
	db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte("shortURL"))
		if err != nil {
			return err
		}
		_, err = tx.CreateBucketIfNotExists([]byte("viewCount"))
		if err != nil {
			return err
		}
		return nil
	})

	ar := &DBHandler{DB: db}

	mux := goji.NewMux()
	mux.HandleFunc(pat.Get("/api"), ar.api)
	mux.HandleFunc(pat.Post("/api"), ar.api)

	mux.HandleFunc(pat.Get("/:shorturl"), ar.gotoshorturl)

	http.ListenAndServe("localhost:"+*port, mux)
}

// 小写hex
func Md5HexFromString(data, salt string) []byte {
	hash := md5.New()
	hash.Write([]byte(data))
	if len(salt) > 0 {
		hash.Write([]byte(salt))
	}
	return []byte(hex.EncodeToString(hash.Sum(nil)))
}

// 生成一个长URL的短链接
func URLshorten(longURL string) []string {
	// 加密字符串
	key := "URLshorten"
	// URL字符表,共62个,下标为0~61
	text := "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
	md5text := Md5HexFromString(longURL, key)
	shortURLs := make([]string, 0, 4)
	for i := 0; i < 4; i++ {
		// 取出MD5中第i个字节,并忽略超过30位的部分
		str := md5text[i*8 : (i+1)*8]
		// 把str里的十六进制表示转化成int
		num, err := strconv.Atoi(fmt.Sprintf("%x", string(str)))
		if err != nil {
			panic(fmt.Errorf("invalid longurl:%v", longURL))
		}
		num &= 0x3FFFFFFF // 选择低30位

		// 取30位的后6位与0x0000003D进行逻辑与操作,结果范围是0~61,作为text的下标选择字符
		// 把num左移5位重复进行,得到6个字符组成短URL
		shortURL := make([]byte, 0, 6)
		for j := 0; j < 6; j++ {
			shortURL = append(shortURL, text[num&0x0000003D])
			num >>= 5
		}
		shortURLs = append(shortURLs, string(shortURL))
	}
	// fmt.Printf("%+v\n", shortURLs)
	// return shortURLs[rand.Intn(len(shortURLs))]
	return shortURLs
}

本文网址: https://golangnote.com/topic/207.html (转载注明出处)
关于GolangNote:记录在工作中使用golang 遇到、面临的相关问题及解决方法。如果你在这里获得一些知识或信息,解决你的编程问题,请考虑捐赠给不幸的人或者你喜欢的慈善机构,除捐赠外,种植树木、志愿服务或减少排碳的行为也很有益处。如果你有任何问题可以在下面 留言
Be the first to comment!
Captcha image
Relative Articles