images.go 23 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028
  1. package server
  2. import (
  3. "bytes"
  4. "crypto/sha256"
  5. "encoding/json"
  6. "errors"
  7. "fmt"
  8. "io"
  9. "io/ioutil"
  10. "log"
  11. "net/http"
  12. "os"
  13. "path"
  14. "path/filepath"
  15. "reflect"
  16. "strconv"
  17. "strings"
  18. "text/template"
  19. "github.com/jmorganca/ollama/api"
  20. "github.com/jmorganca/ollama/parser"
  21. )
  22. type RegistryOptions struct {
  23. Insecure bool
  24. Username string
  25. Password string
  26. }
  27. type Model struct {
  28. Name string `json:"name"`
  29. ModelPath string
  30. Template string
  31. System string
  32. Options api.Options
  33. }
  34. func (m *Model) Prompt(request api.GenerateRequest) (string, error) {
  35. tmpl, err := template.New("").Parse(m.Template)
  36. if err != nil {
  37. return "", err
  38. }
  39. var vars struct {
  40. First bool
  41. System string
  42. Prompt string
  43. // deprecated: versions <= 0.0.7 used this to omit the system prompt
  44. Context []int
  45. }
  46. vars.First = len(request.Context) == 0
  47. vars.System = m.System
  48. vars.Prompt = request.Prompt
  49. vars.Context = request.Context
  50. var sb strings.Builder
  51. if err := tmpl.Execute(&sb, vars); err != nil {
  52. return "", err
  53. }
  54. return sb.String(), nil
  55. }
  56. type ManifestV2 struct {
  57. SchemaVersion int `json:"schemaVersion"`
  58. MediaType string `json:"mediaType"`
  59. Config Layer `json:"config"`
  60. Layers []*Layer `json:"layers"`
  61. }
  62. type Layer struct {
  63. MediaType string `json:"mediaType"`
  64. Digest string `json:"digest"`
  65. Size int `json:"size"`
  66. }
  67. type LayerReader struct {
  68. Layer
  69. io.Reader
  70. }
  71. type ConfigV2 struct {
  72. Architecture string `json:"architecture"`
  73. OS string `json:"os"`
  74. RootFS RootFS `json:"rootfs"`
  75. }
  76. type RootFS struct {
  77. Type string `json:"type"`
  78. DiffIDs []string `json:"diff_ids"`
  79. }
  80. func (m *ManifestV2) GetTotalSize() int {
  81. var total int
  82. for _, layer := range m.Layers {
  83. total += layer.Size
  84. }
  85. total += m.Config.Size
  86. return total
  87. }
  88. func GetManifest(mp ModelPath) (*ManifestV2, error) {
  89. fp, err := mp.GetManifestPath(false)
  90. if err != nil {
  91. return nil, err
  92. }
  93. if _, err = os.Stat(fp); err != nil {
  94. return nil, err
  95. }
  96. var manifest *ManifestV2
  97. bts, err := os.ReadFile(fp)
  98. if err != nil {
  99. return nil, fmt.Errorf("couldn't open file '%s'", fp)
  100. }
  101. if err := json.Unmarshal(bts, &manifest); err != nil {
  102. return nil, err
  103. }
  104. return manifest, nil
  105. }
  106. func GetModel(name string) (*Model, error) {
  107. mp := ParseModelPath(name)
  108. manifest, err := GetManifest(mp)
  109. if err != nil {
  110. return nil, err
  111. }
  112. model := &Model{
  113. Name: mp.GetFullTagname(),
  114. }
  115. for _, layer := range manifest.Layers {
  116. filename, err := GetBlobsPath(layer.Digest)
  117. if err != nil {
  118. return nil, err
  119. }
  120. switch layer.MediaType {
  121. case "application/vnd.ollama.image.model":
  122. model.ModelPath = filename
  123. case "application/vnd.ollama.image.template":
  124. bts, err := os.ReadFile(filename)
  125. if err != nil {
  126. return nil, err
  127. }
  128. model.Template = string(bts)
  129. case "application/vnd.ollama.image.system":
  130. bts, err := os.ReadFile(filename)
  131. if err != nil {
  132. return nil, err
  133. }
  134. model.System = string(bts)
  135. case "application/vnd.ollama.image.prompt":
  136. bts, err := os.ReadFile(filename)
  137. if err != nil {
  138. return nil, err
  139. }
  140. model.Template = string(bts)
  141. case "application/vnd.ollama.image.params":
  142. params, err := os.Open(filename)
  143. if err != nil {
  144. return nil, err
  145. }
  146. defer params.Close()
  147. var opts api.Options
  148. if err = json.NewDecoder(params).Decode(&opts); err != nil {
  149. return nil, err
  150. }
  151. model.Options = opts
  152. }
  153. }
  154. return model, nil
  155. }
  156. func CreateModel(name string, path string, fn func(status string)) error {
  157. mf, err := os.Open(path)
  158. if err != nil {
  159. fn(fmt.Sprintf("couldn't open modelfile '%s'", path))
  160. return fmt.Errorf("failed to open file: %w", err)
  161. }
  162. defer mf.Close()
  163. fn("parsing modelfile")
  164. commands, err := parser.Parse(mf)
  165. if err != nil {
  166. return err
  167. }
  168. var layers []*LayerReader
  169. params := make(map[string]string)
  170. for _, c := range commands {
  171. log.Printf("[%s] - %s\n", c.Name, c.Args)
  172. switch c.Name {
  173. case "model":
  174. fn("looking for model")
  175. mf, err := GetManifest(ParseModelPath(c.Args))
  176. if err != nil {
  177. fp := c.Args
  178. // If filePath starts with ~/, replace it with the user's home directory.
  179. if strings.HasPrefix(fp, "~/") {
  180. parts := strings.Split(fp, "/")
  181. home, err := os.UserHomeDir()
  182. if err != nil {
  183. return fmt.Errorf("failed to open file: %v", err)
  184. }
  185. fp = filepath.Join(home, filepath.Join(parts[1:]...))
  186. }
  187. // If filePath is not an absolute path, make it relative to the modelfile path
  188. if !filepath.IsAbs(fp) {
  189. fp = filepath.Join(filepath.Dir(path), fp)
  190. }
  191. fn("creating model layer")
  192. file, err := os.Open(fp)
  193. if err != nil {
  194. return fmt.Errorf("failed to open file: %v", err)
  195. }
  196. defer file.Close()
  197. l, err := CreateLayer(file)
  198. if err != nil {
  199. return fmt.Errorf("failed to create layer: %v", err)
  200. }
  201. l.MediaType = "application/vnd.ollama.image.model"
  202. layers = append(layers, l)
  203. } else {
  204. log.Printf("manifest = %#v", mf)
  205. for _, l := range mf.Layers {
  206. newLayer, err := GetLayerWithBufferFromLayer(l)
  207. if err != nil {
  208. return err
  209. }
  210. layers = append(layers, newLayer)
  211. }
  212. }
  213. case "license", "template", "system", "prompt":
  214. fn(fmt.Sprintf("creating %s layer", c.Name))
  215. // remove the prompt layer if one exists
  216. mediaType := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)
  217. layers = removeLayerFromLayers(layers, mediaType)
  218. layer, err := CreateLayer(strings.NewReader(c.Args))
  219. if err != nil {
  220. return err
  221. }
  222. layer.MediaType = mediaType
  223. layers = append(layers, layer)
  224. default:
  225. params[c.Name] = c.Args
  226. }
  227. }
  228. // Create a single layer for the parameters
  229. if len(params) > 0 {
  230. fn("creating parameter layer")
  231. layers = removeLayerFromLayers(layers, "application/vnd.ollama.image.params")
  232. paramData, err := paramsToReader(params)
  233. if err != nil {
  234. return fmt.Errorf("couldn't create params json: %v", err)
  235. }
  236. l, err := CreateLayer(paramData)
  237. if err != nil {
  238. return fmt.Errorf("failed to create layer: %v", err)
  239. }
  240. l.MediaType = "application/vnd.ollama.image.params"
  241. layers = append(layers, l)
  242. }
  243. digests, err := getLayerDigests(layers)
  244. if err != nil {
  245. return err
  246. }
  247. var manifestLayers []*Layer
  248. for _, l := range layers {
  249. manifestLayers = append(manifestLayers, &l.Layer)
  250. }
  251. // Create a layer for the config object
  252. fn("creating config layer")
  253. cfg, err := createConfigLayer(digests)
  254. if err != nil {
  255. return err
  256. }
  257. layers = append(layers, cfg)
  258. err = SaveLayers(layers, fn, false)
  259. if err != nil {
  260. return err
  261. }
  262. // Create the manifest
  263. fn("writing manifest")
  264. err = CreateManifest(name, cfg, manifestLayers)
  265. if err != nil {
  266. return err
  267. }
  268. fn("success")
  269. return nil
  270. }
  271. func removeLayerFromLayers(layers []*LayerReader, mediaType string) []*LayerReader {
  272. j := 0
  273. for _, l := range layers {
  274. if l.MediaType != mediaType {
  275. layers[j] = l
  276. j++
  277. }
  278. }
  279. return layers[:j]
  280. }
  281. func SaveLayers(layers []*LayerReader, fn func(status string), force bool) error {
  282. // Write each of the layers to disk
  283. for _, layer := range layers {
  284. fp, err := GetBlobsPath(layer.Digest)
  285. if err != nil {
  286. return err
  287. }
  288. _, err = os.Stat(fp)
  289. if os.IsNotExist(err) || force {
  290. fn(fmt.Sprintf("writing layer %s", layer.Digest))
  291. out, err := os.Create(fp)
  292. if err != nil {
  293. log.Printf("couldn't create %s", fp)
  294. return err
  295. }
  296. defer out.Close()
  297. if _, err = io.Copy(out, layer.Reader); err != nil {
  298. return err
  299. }
  300. } else {
  301. fn(fmt.Sprintf("using already created layer %s", layer.Digest))
  302. }
  303. }
  304. return nil
  305. }
  306. func CreateManifest(name string, cfg *LayerReader, layers []*Layer) error {
  307. mp := ParseModelPath(name)
  308. manifest := ManifestV2{
  309. SchemaVersion: 2,
  310. MediaType: "application/vnd.docker.distribution.manifest.v2+json",
  311. Config: Layer{
  312. MediaType: cfg.MediaType,
  313. Size: cfg.Size,
  314. Digest: cfg.Digest,
  315. },
  316. Layers: layers,
  317. }
  318. manifestJSON, err := json.Marshal(manifest)
  319. if err != nil {
  320. return err
  321. }
  322. fp, err := mp.GetManifestPath(true)
  323. if err != nil {
  324. return err
  325. }
  326. return os.WriteFile(fp, manifestJSON, 0o644)
  327. }
  328. func GetLayerWithBufferFromLayer(layer *Layer) (*LayerReader, error) {
  329. fp, err := GetBlobsPath(layer.Digest)
  330. if err != nil {
  331. return nil, err
  332. }
  333. file, err := os.Open(fp)
  334. if err != nil {
  335. return nil, fmt.Errorf("could not open blob: %w", err)
  336. }
  337. defer file.Close()
  338. newLayer, err := CreateLayer(file)
  339. if err != nil {
  340. return nil, err
  341. }
  342. newLayer.MediaType = layer.MediaType
  343. return newLayer, nil
  344. }
  345. func paramsToReader(params map[string]string) (io.ReadSeeker, error) {
  346. opts := api.DefaultOptions()
  347. typeOpts := reflect.TypeOf(opts)
  348. // build map of json struct tags
  349. jsonOpts := make(map[string]reflect.StructField)
  350. for _, field := range reflect.VisibleFields(typeOpts) {
  351. jsonTag := strings.Split(field.Tag.Get("json"), ",")[0]
  352. if jsonTag != "" {
  353. jsonOpts[jsonTag] = field
  354. }
  355. }
  356. valueOpts := reflect.ValueOf(&opts).Elem()
  357. // iterate params and set values based on json struct tags
  358. for key, val := range params {
  359. if opt, ok := jsonOpts[key]; ok {
  360. field := valueOpts.FieldByName(opt.Name)
  361. if field.IsValid() && field.CanSet() {
  362. switch field.Kind() {
  363. case reflect.Float32:
  364. floatVal, err := strconv.ParseFloat(val, 32)
  365. if err != nil {
  366. return nil, fmt.Errorf("invalid float value %s", val)
  367. }
  368. field.SetFloat(floatVal)
  369. case reflect.Int:
  370. intVal, err := strconv.ParseInt(val, 10, 0)
  371. if err != nil {
  372. return nil, fmt.Errorf("invalid int value %s", val)
  373. }
  374. field.SetInt(intVal)
  375. case reflect.Bool:
  376. boolVal, err := strconv.ParseBool(val)
  377. if err != nil {
  378. return nil, fmt.Errorf("invalid bool value %s", val)
  379. }
  380. field.SetBool(boolVal)
  381. case reflect.String:
  382. field.SetString(val)
  383. default:
  384. return nil, fmt.Errorf("unknown type %s for %s", field.Kind(), key)
  385. }
  386. }
  387. }
  388. }
  389. bts, err := json.Marshal(opts)
  390. if err != nil {
  391. return nil, err
  392. }
  393. return bytes.NewReader(bts), nil
  394. }
  395. func getLayerDigests(layers []*LayerReader) ([]string, error) {
  396. var digests []string
  397. for _, l := range layers {
  398. if l.Digest == "" {
  399. return nil, fmt.Errorf("layer is missing a digest")
  400. }
  401. digests = append(digests, l.Digest)
  402. }
  403. return digests, nil
  404. }
  405. // CreateLayer creates a Layer object from a given file
  406. func CreateLayer(f io.ReadSeeker) (*LayerReader, error) {
  407. digest, size := GetSHA256Digest(f)
  408. f.Seek(0, 0)
  409. layer := &LayerReader{
  410. Layer: Layer{
  411. MediaType: "application/vnd.docker.image.rootfs.diff.tar",
  412. Digest: digest,
  413. Size: size,
  414. },
  415. Reader: f,
  416. }
  417. return layer, nil
  418. }
  419. func DeleteModel(name string) error {
  420. mp := ParseModelPath(name)
  421. manifest, err := GetManifest(mp)
  422. if err != nil {
  423. return err
  424. }
  425. deleteMap := make(map[string]bool)
  426. for _, layer := range manifest.Layers {
  427. deleteMap[layer.Digest] = true
  428. }
  429. deleteMap[manifest.Config.Digest] = true
  430. fp, err := GetManifestPath()
  431. if err != nil {
  432. return err
  433. }
  434. err = filepath.Walk(fp, func(path string, info os.FileInfo, err error) error {
  435. if err != nil {
  436. return err
  437. }
  438. if !info.IsDir() {
  439. path := path[len(fp)+1:]
  440. slashIndex := strings.LastIndex(path, "/")
  441. if slashIndex == -1 {
  442. return nil
  443. }
  444. tag := path[:slashIndex] + ":" + path[slashIndex+1:]
  445. fmp := ParseModelPath(tag)
  446. // skip the manifest we're trying to delete
  447. if mp.GetFullTagname() == fmp.GetFullTagname() {
  448. return nil
  449. }
  450. // save (i.e. delete from the deleteMap) any files used in other manifests
  451. manifest, err := GetManifest(fmp)
  452. if err != nil {
  453. log.Printf("skipping file: %s", fp)
  454. return nil
  455. }
  456. for _, layer := range manifest.Layers {
  457. delete(deleteMap, layer.Digest)
  458. }
  459. delete(deleteMap, manifest.Config.Digest)
  460. }
  461. return nil
  462. })
  463. // only delete the files which are still in the deleteMap
  464. for k, v := range deleteMap {
  465. if v {
  466. fp, err := GetBlobsPath(k)
  467. if err != nil {
  468. log.Printf("couldn't get file path for '%s': %v", k, err)
  469. continue
  470. }
  471. if err := os.Remove(fp); err != nil {
  472. log.Printf("couldn't remove file '%s': %v", fp, err)
  473. continue
  474. }
  475. }
  476. }
  477. fp, err = mp.GetManifestPath(false)
  478. if err != nil {
  479. return err
  480. }
  481. err = os.Remove(fp)
  482. if err != nil {
  483. log.Printf("couldn't remove manifest file '%s': %v", fp, err)
  484. return err
  485. }
  486. return nil
  487. }
  488. func PushModel(name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
  489. mp := ParseModelPath(name)
  490. fn(api.ProgressResponse{Status: "retrieving manifest"})
  491. manifest, err := GetManifest(mp)
  492. if err != nil {
  493. fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
  494. return err
  495. }
  496. var layers []*Layer
  497. var total int
  498. var completed int
  499. for _, layer := range manifest.Layers {
  500. layers = append(layers, layer)
  501. total += layer.Size
  502. }
  503. layers = append(layers, &manifest.Config)
  504. total += manifest.Config.Size
  505. for _, layer := range layers {
  506. exists, err := checkBlobExistence(mp, layer.Digest, regOpts)
  507. if err != nil {
  508. return err
  509. }
  510. if exists {
  511. completed += layer.Size
  512. fn(api.ProgressResponse{
  513. Status: "using existing layer",
  514. Digest: layer.Digest,
  515. Total: total,
  516. Completed: completed,
  517. })
  518. continue
  519. }
  520. fn(api.ProgressResponse{
  521. Status: "starting upload",
  522. Digest: layer.Digest,
  523. Total: total,
  524. Completed: completed,
  525. })
  526. location, err := startUpload(mp, regOpts)
  527. if err != nil {
  528. log.Printf("couldn't start upload: %v", err)
  529. return err
  530. }
  531. err = uploadBlob(location, layer, regOpts)
  532. if err != nil {
  533. log.Printf("error uploading blob: %v", err)
  534. return err
  535. }
  536. completed += layer.Size
  537. fn(api.ProgressResponse{
  538. Status: "upload complete",
  539. Digest: layer.Digest,
  540. Total: total,
  541. Completed: completed,
  542. })
  543. }
  544. fn(api.ProgressResponse{
  545. Status: "pushing manifest",
  546. Total: total,
  547. Completed: completed,
  548. })
  549. url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
  550. headers := map[string]string{
  551. "Content-Type": "application/vnd.docker.distribution.manifest.v2+json",
  552. }
  553. manifestJSON, err := json.Marshal(manifest)
  554. if err != nil {
  555. return err
  556. }
  557. resp, err := makeRequest("PUT", url, headers, bytes.NewReader(manifestJSON), regOpts)
  558. if err != nil {
  559. return err
  560. }
  561. defer resp.Body.Close()
  562. // Check for success: For a successful upload, the Docker registry will respond with a 201 Created
  563. if resp.StatusCode != http.StatusCreated {
  564. body, _ := io.ReadAll(resp.Body)
  565. return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body))
  566. }
  567. fn(api.ProgressResponse{
  568. Status: "success",
  569. Total: total,
  570. Completed: completed,
  571. })
  572. return nil
  573. }
  574. func PullModel(name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
  575. mp := ParseModelPath(name)
  576. fn(api.ProgressResponse{Status: "pulling manifest"})
  577. manifest, err := pullModelManifest(mp, regOpts)
  578. if err != nil {
  579. return fmt.Errorf("pull model manifest: %q", err)
  580. }
  581. var layers []*Layer
  582. layers = append(layers, manifest.Layers...)
  583. layers = append(layers, &manifest.Config)
  584. for _, layer := range layers {
  585. if err := downloadBlob(mp, layer.Digest, regOpts, fn); err != nil {
  586. return err
  587. }
  588. }
  589. fn(api.ProgressResponse{Status: "verifying sha256 digest"})
  590. for _, layer := range layers {
  591. if err := verifyBlob(layer.Digest); err != nil {
  592. return err
  593. }
  594. }
  595. fn(api.ProgressResponse{Status: "writing manifest"})
  596. manifestJSON, err := json.Marshal(manifest)
  597. if err != nil {
  598. return err
  599. }
  600. fp, err := mp.GetManifestPath(true)
  601. if err != nil {
  602. return err
  603. }
  604. err = os.WriteFile(fp, manifestJSON, 0o644)
  605. if err != nil {
  606. log.Printf("couldn't write to %s", fp)
  607. return err
  608. }
  609. fn(api.ProgressResponse{Status: "success"})
  610. return nil
  611. }
  612. func pullModelManifest(mp ModelPath, regOpts *RegistryOptions) (*ManifestV2, error) {
  613. url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
  614. headers := map[string]string{
  615. "Accept": "application/vnd.docker.distribution.manifest.v2+json",
  616. }
  617. resp, err := makeRequest("GET", url, headers, nil, regOpts)
  618. if err != nil {
  619. log.Printf("couldn't get manifest: %v", err)
  620. return nil, err
  621. }
  622. defer resp.Body.Close()
  623. // Check for success: For a successful upload, the Docker registry will respond with a 201 Created
  624. if resp.StatusCode != http.StatusOK {
  625. body, _ := io.ReadAll(resp.Body)
  626. return nil, fmt.Errorf("registry responded with code %d: %s", resp.StatusCode, body)
  627. }
  628. var m *ManifestV2
  629. if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
  630. return nil, err
  631. }
  632. return m, err
  633. }
  634. func createConfigLayer(layers []string) (*LayerReader, error) {
  635. // TODO change architecture and OS
  636. config := ConfigV2{
  637. Architecture: "arm64",
  638. OS: "linux",
  639. RootFS: RootFS{
  640. Type: "layers",
  641. DiffIDs: layers,
  642. },
  643. }
  644. configJSON, err := json.Marshal(config)
  645. if err != nil {
  646. return nil, err
  647. }
  648. digest, size := GetSHA256Digest(bytes.NewBuffer(configJSON))
  649. layer := &LayerReader{
  650. Layer: Layer{
  651. MediaType: "application/vnd.docker.container.image.v1+json",
  652. Digest: digest,
  653. Size: size,
  654. },
  655. Reader: bytes.NewBuffer(configJSON),
  656. }
  657. return layer, nil
  658. }
  659. // GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
  660. func GetSHA256Digest(r io.Reader) (string, int) {
  661. h := sha256.New()
  662. n, err := io.Copy(h, r)
  663. if err != nil {
  664. log.Fatal(err)
  665. }
  666. return fmt.Sprintf("sha256:%x", h.Sum(nil)), int(n)
  667. }
  668. func startUpload(mp ModelPath, regOpts *RegistryOptions) (string, error) {
  669. url := fmt.Sprintf("%s/v2/%s/blobs/uploads/", mp.Registry, mp.GetNamespaceRepository())
  670. resp, err := makeRequest("POST", url, nil, nil, regOpts)
  671. if err != nil {
  672. log.Printf("couldn't start upload: %v", err)
  673. return "", err
  674. }
  675. defer resp.Body.Close()
  676. // Check for success
  677. if resp.StatusCode != http.StatusAccepted {
  678. body, _ := io.ReadAll(resp.Body)
  679. return "", fmt.Errorf("registry responded with code %d: %s", resp.StatusCode, body)
  680. }
  681. // Extract UUID location from header
  682. location := resp.Header.Get("Location")
  683. if location == "" {
  684. return "", fmt.Errorf("location header is missing in response")
  685. }
  686. return location, nil
  687. }
  688. // Function to check if a blob already exists in the Docker registry
  689. func checkBlobExistence(mp ModelPath, digest string, regOpts *RegistryOptions) (bool, error) {
  690. url := fmt.Sprintf("%s/v2/%s/blobs/%s", mp.Registry, mp.GetNamespaceRepository(), digest)
  691. resp, err := makeRequest("HEAD", url, nil, nil, regOpts)
  692. if err != nil {
  693. log.Printf("couldn't check for blob: %v", err)
  694. return false, err
  695. }
  696. defer resp.Body.Close()
  697. // Check for success: If the blob exists, the Docker registry will respond with a 200 OK
  698. return resp.StatusCode == http.StatusOK, nil
  699. }
  700. func uploadBlob(location string, layer *Layer, regOpts *RegistryOptions) error {
  701. // Create URL
  702. url := fmt.Sprintf("%s&digest=%s", location, layer.Digest)
  703. headers := make(map[string]string)
  704. headers["Content-Length"] = fmt.Sprintf("%d", layer.Size)
  705. headers["Content-Type"] = "application/octet-stream"
  706. // TODO change from monolithic uploads to chunked uploads
  707. // TODO allow resumability
  708. // TODO allow canceling uploads via DELETE
  709. // TODO allow cross repo blob mount
  710. fp, err := GetBlobsPath(layer.Digest)
  711. if err != nil {
  712. return err
  713. }
  714. f, err := os.Open(fp)
  715. if err != nil {
  716. return err
  717. }
  718. resp, err := makeRequest("PUT", url, headers, f, regOpts)
  719. if err != nil {
  720. log.Printf("couldn't upload blob: %v", err)
  721. return err
  722. }
  723. defer resp.Body.Close()
  724. // Check for success: For a successful upload, the Docker registry will respond with a 201 Created
  725. if resp.StatusCode != http.StatusCreated {
  726. body, _ := io.ReadAll(resp.Body)
  727. return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body))
  728. }
  729. return nil
  730. }
  731. func downloadBlob(mp ModelPath, digest string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
  732. fp, err := GetBlobsPath(digest)
  733. if err != nil {
  734. return err
  735. }
  736. if fi, _ := os.Stat(fp); fi != nil {
  737. // we already have the file, so return
  738. fn(api.ProgressResponse{
  739. Digest: digest,
  740. Total: int(fi.Size()),
  741. Completed: int(fi.Size()),
  742. })
  743. return nil
  744. }
  745. var size int64
  746. fi, err := os.Stat(fp + "-partial")
  747. switch {
  748. case errors.Is(err, os.ErrNotExist):
  749. // noop, file doesn't exist so create it
  750. case err != nil:
  751. return fmt.Errorf("stat: %w", err)
  752. default:
  753. size = fi.Size()
  754. }
  755. url := fmt.Sprintf("%s/v2/%s/blobs/%s", mp.Registry, mp.GetNamespaceRepository(), digest)
  756. headers := map[string]string{
  757. "Range": fmt.Sprintf("bytes=%d-", size),
  758. }
  759. resp, err := makeRequest("GET", url, headers, nil, regOpts)
  760. if err != nil {
  761. log.Printf("couldn't download blob: %v", err)
  762. return err
  763. }
  764. defer resp.Body.Close()
  765. if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent {
  766. body, _ := ioutil.ReadAll(resp.Body)
  767. return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body))
  768. }
  769. err = os.MkdirAll(path.Dir(fp), 0o700)
  770. if err != nil {
  771. return fmt.Errorf("make blobs directory: %w", err)
  772. }
  773. out, err := os.OpenFile(fp+"-partial", os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644)
  774. if err != nil {
  775. panic(err)
  776. }
  777. defer out.Close()
  778. remaining, _ := strconv.ParseInt(resp.Header.Get("Content-Length"), 10, 64)
  779. completed := size
  780. total := remaining + completed
  781. for {
  782. fn(api.ProgressResponse{
  783. Status: fmt.Sprintf("downloading %s", digest),
  784. Digest: digest,
  785. Total: int(total),
  786. Completed: int(completed),
  787. })
  788. if completed >= total {
  789. if err := out.Close(); err != nil {
  790. return err
  791. }
  792. if err := os.Rename(fp+"-partial", fp); err != nil {
  793. fn(api.ProgressResponse{
  794. Status: fmt.Sprintf("error renaming file: %v", err),
  795. Digest: digest,
  796. Total: int(total),
  797. Completed: int(completed),
  798. })
  799. return err
  800. }
  801. break
  802. }
  803. n, err := io.CopyN(out, resp.Body, 8192)
  804. if err != nil && !errors.Is(err, io.EOF) {
  805. return err
  806. }
  807. completed += n
  808. }
  809. log.Printf("success getting %s\n", digest)
  810. return nil
  811. }
  812. func makeRequest(method, url string, headers map[string]string, body io.Reader, regOpts *RegistryOptions) (*http.Response, error) {
  813. if !strings.HasPrefix(url, "http") {
  814. if regOpts.Insecure {
  815. url = "http://" + url
  816. } else {
  817. url = "https://" + url
  818. }
  819. }
  820. log.Printf("url = %s", url)
  821. req, err := http.NewRequest(method, url, body)
  822. if err != nil {
  823. return nil, err
  824. }
  825. for k, v := range headers {
  826. req.Header.Set(k, v)
  827. }
  828. // TODO: better auth
  829. if regOpts.Username != "" && regOpts.Password != "" {
  830. req.SetBasicAuth(regOpts.Username, regOpts.Password)
  831. }
  832. client := &http.Client{
  833. CheckRedirect: func(req *http.Request, via []*http.Request) error {
  834. if len(via) >= 10 {
  835. return fmt.Errorf("too many redirects")
  836. }
  837. log.Printf("redirected to: %s\n", req.URL)
  838. return nil
  839. },
  840. }
  841. resp, err := client.Do(req)
  842. if err != nil {
  843. return nil, err
  844. }
  845. return resp, nil
  846. }
  847. func verifyBlob(digest string) error {
  848. fp, err := GetBlobsPath(digest)
  849. if err != nil {
  850. return err
  851. }
  852. f, err := os.Open(fp)
  853. if err != nil {
  854. return err
  855. }
  856. defer f.Close()
  857. fileDigest, _ := GetSHA256Digest(f)
  858. if digest != fileDigest {
  859. return fmt.Errorf("digest mismatch: want %s, got %s", digest, fileDigest)
  860. }
  861. return nil
  862. }