package main import ( "database/sql" "flag" "fmt" "log" "math/rand" "net/http" "net/url" "os" "regexp" "strings" "time" "github.com/4rkal/shortr/models" "github.com/4rkal/shortr/views" "github.com/labstack/echo/v4" "github.com/labstack/echo/v4/middleware" _ "github.com/lib/pq" ) const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" type FormData struct { Url string `form:"url"` } type StatsFormData struct { Id string `form:"id"` } var linkMap = map[string]*models.Link{} var baseurl *string var db *sql.DB func init() { baseurl = flag.String("url", "127.0.0.1:8080", "The URL (domain) that the server is running on") flag.Parse() var err error db, err = sql.Open("postgres", os.Getenv("DATABASE_URL")) if err != nil { panic(fmt.Sprintf("Failed to connect to database: %v", err)) } if err := db.Ping(); err != nil { panic(fmt.Sprintf("Failed to ping database: %v", err)) } if err := loadLinksIntoMemory(); err != nil { panic(fmt.Sprintf("Failed to load links into memory: %v", err)) } _, err = db.Exec(` CREATE TABLE IF NOT EXISTS links ( id VARCHAR(255) PRIMARY KEY, url TEXT NOT NULL, clicks INTEGER NOT NULL DEFAULT 0 ) `) if err != nil { log.Fatal("Error creating table: ", err) } if err := loadLinksIntoMemory(); err != nil { log.Fatalf("Failed to load links into memory: %v", err) } } func main() { e := echo.New() e.Use(middleware.Logger()) e.Use(middleware.Recover()) e.Use(middleware.Secure()) e.GET("/stats", StatsHandler) e.GET("/:id", RedirectHandler) e.GET("/:id/", RedirectHandler) e.GET("/", IndexHandler) e.POST("/submit", SubmitHandler) e.GET("/health", HealthHandler) e.Logger.Fatal(e.Start(":8080")) } func RedirectHandler(c echo.Context) error { id := c.Param("id") link, found := linkMap[id] if !found { return c.String(http.StatusNotFound, "Link not found") } if !strings.Contains(link.Url, "://") { link.Url = "http://" + link.Url } link.Clicks = link.Clicks + 1 go func(id string, clicks int) { _, err := db.Exec("UPDATE links SET clicks = $1 WHERE id = $2", clicks, id) if err != nil { fmt.Printf("Failed to update clicks for id %s: %v\n", id, err) } }(id, link.Clicks) return c.Redirect(http.StatusMovedPermanently, link.Url) } func IndexHandler(c echo.Context) error { return views.Index().Render(c.Request().Context(), c.Response()) } func SubmitHandler(c echo.Context) error { var data FormData if err := c.Bind(&data); err != nil { return err } fmt.Println(data) if !isURL(data.Url) { return c.JSON(http.StatusBadRequest, "not a valid url") } var id string for { id = generateRandomString(6) if _, exists := linkMap[id]; !exists { break } } linkMap[id] = &models.Link{Id: id, Url: data.Url} _, err := db.Exec("INSERT INTO links (id, url, clicks) VALUES ($1, $2, $3)", id, data.Url, 0) if err != nil { return fmt.Errorf("failed to save link to database: %v", err) } return views.Submission(id, *baseurl).Render(c.Request().Context(), c.Response()) } func StatsHandler(c echo.Context) error { id := c.QueryParam("id") if id == "" { return views.StatsForm().Render(c.Request().Context(), c.Response()) } else { link, found := linkMap[id] if !found { return c.String(http.StatusNotFound, "Id not found") } return views.Stats(link).Render(c.Request().Context(), c.Response()) } } func HealthHandler(c echo.Context) error { return c.JSON(http.StatusOK, "ok") } func isURL(s string) bool { if !strings.Contains(s, "://") { s = "http://" + s } parsedUrl, err := url.ParseRequestURI(s) if err != nil { return false } if parsedUrl.Host == "" { return false } domainRegex := `^([a-zA-Z0-9-]+\.)+[a-zA-Z]{2,}$` matched, err := regexp.MatchString(domainRegex, parsedUrl.Host) if err != nil || !matched { return false } return true } func generateRandomString(length int) string { seededRand := rand.New(rand.NewSource(time.Now().UnixNano())) var result []byte for i := 0; i < length; i++ { index := seededRand.Intn(len(charset)) result = append(result, charset[index]) } return string(result) } func loadLinksIntoMemory() error { rows, err := db.Query("SELECT id, url, clicks FROM links") if err != nil { return fmt.Errorf("query error: %v", err) } defer rows.Close() for rows.Next() { var id, url string var clicks int if err := rows.Scan(&id, &url, &clicks); err != nil { return fmt.Errorf("scan error: %v", err) } linkMap[id] = &models.Link{Id: id, Url: url, Clicks: clicks} } return nil }